diff --git a/README.md b/README.md index 38d5a76c4..d1b115bb1 100644 --- a/README.md +++ b/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. diff --git a/docs/usage.rst b/docs/usage.rst index 44f176c7f..fe33de503 100644 --- a/docs/usage.rst +++ b/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", diff --git a/examples/basic_example.ipynb b/examples/basic_example.ipynb index 44acdfac2..8ebf16c78 100644 --- a/examples/basic_example.ipynb +++ b/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 = ''\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" ] }, { diff --git a/examples/configs/estimators/default_claim_estimators.yaml b/examples/configs/estimators/default_claim_estimators.yaml index 50cd842ee..019bfeddb 100644 --- a/examples/configs/estimators/default_claim_estimators.yaml +++ b/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 \ No newline at end of file +- name: TokenSARClaim diff --git a/examples/configs/estimators/default_estimators_blackbox.yaml b/examples/configs/estimators/default_estimators_blackbox.yaml new file mode 100644 index 000000000..268e9536e --- /dev/null +++ b/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" diff --git a/examples/configs/estimators/default_estimators_greybox.yaml b/examples/configs/estimators/default_estimators_greybox.yaml new file mode 100644 index 000000000..253e1a709 --- /dev/null +++ b/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 diff --git a/examples/configs/model/gpt_4.1.yaml b/examples/configs/model/gpt_4.1.yaml new file mode 100644 index 000000000..07d80ee6d --- /dev/null +++ b/examples/configs/model/gpt_4.1.yaml @@ -0,0 +1,3 @@ +path: gpt-4.1 +type: Blackbox +provider: openai diff --git a/examples/configs/model/llama_huggingface.yaml b/examples/configs/model/llama_huggingface.yaml new file mode 100644 index 000000000..c425b41c2 --- /dev/null +++ b/examples/configs/model/llama_huggingface.yaml @@ -0,0 +1,3 @@ +path: meta-llama/Meta-Llama-3.1-8B-Instruct +type: Blackbox +provider: huggingface diff --git a/examples/configs/model/llama_together.yaml b/examples/configs/model/llama_together.yaml new file mode 100644 index 000000000..38b054810 --- /dev/null +++ b/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 diff --git a/examples/other/ats_example.ipynb b/examples/other/ats_example.ipynb index 6c51f1dd6..8f9d7e6bf 100644 --- a/examples/other/ats_example.ipynb +++ b/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", diff --git a/examples/other/claim_level_example_ar.ipynb b/examples/other/claim_level_example_ar.ipynb index aaafa5c8a..a3c9bbdcc 100644 --- a/examples/other/claim_level_example_ar.ipynb +++ b/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", diff --git a/examples/other/claim_level_example_en.ipynb b/examples/other/claim_level_example_en.ipynb index 41cc2d4cf..a6d7b781e 100644 --- a/examples/other/claim_level_example_en.ipynb +++ b/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", diff --git a/examples/other/claim_level_example_ru.ipynb b/examples/other/claim_level_example_ru.ipynb index 094a84bc3..11b28e2ff 100644 --- a/examples/other/claim_level_example_ru.ipynb +++ b/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", diff --git a/examples/other/claim_level_example_zh.ipynb b/examples/other/claim_level_example_zh.ipynb index adf4a1686..ce95b2f43 100644 --- a/examples/other/claim_level_example_zh.ipynb +++ b/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", diff --git a/examples/other/mt_example.ipynb b/examples/other/mt_example.ipynb index dcca5f3ee..f5fb3bdc1 100644 --- a/examples/other/mt_example.ipynb +++ b/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", diff --git a/examples/other/qa_example.ipynb b/examples/other/qa_example.ipynb index 9010a1ad8..c97c31f82 100644 --- a/examples/other/qa_example.ipynb +++ b/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", diff --git a/requirements.txt b/requirements.txt index 3588eab92..12c7c9244 100644 --- a/requirements.txt +++ b/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 diff --git a/scripts/polygraph_eval b/scripts/polygraph_eval index 173223bf0..21c85dc22 100755 --- a/scripts/polygraph_eval +++ b/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: diff --git a/src/lm_polygraph/__init__.py b/src/lm_polygraph/__init__.py index 7330308d8..0645c0304 100644 --- a/src/lm_polygraph/__init__.py +++ b/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 diff --git a/src/lm_polygraph/defaults/register_default_stat_calculators.py b/src/lm_polygraph/defaults/register_default_stat_calculators.py index b7fa94fc7..94b723ef1 100644 --- a/src/lm_polygraph/defaults/register_default_stat_calculators.py +++ b/src/lm_polygraph/defaults/register_default_stat_calculators.py @@ -1,4 +1,4 @@ -from typing import List, Optional +from typing import Any, List, Optional from omegaconf import OmegaConf from pathlib import Path @@ -12,14 +12,23 @@ 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, + model: Optional[Any] = None, ) -> List[StatCalculatorContainer]: """ Specifies the list of the default stat_calculators that could be used in the evaluation scripts and estimate_uncertainty() function with default configurations. + + Args: + model_type: Type of the model ("Blackbox", "Whitebox", or "VisualLM"). + language: Dataset language code. + hf_cache: Optional Hugging Face cache location. + output_attentions: Whether whitebox models should output attentions. + output_hidden_states: Whether whitebox models should output hidden states. + deberta_batch_size: Batch size for NLI-based calculators. + model: Optional model instance used to derive provider-specific behavior (e.g., logprob support). """ all_stat_calculators = [] @@ -74,19 +83,34 @@ def _register( {"nli_model": nli_model_cfg}, ) _register(SemanticClassesCalculator) + _register( + SequenceCrossEncoderSimilarityMatrixCalculator, + "lm_polygraph.defaults.stat_calculator_builders.default_SequenceCrossEncoderSimilarityMatrixCalculator", + { + "batch_size": deberta_batch_size, + "cross_encoder_name": "cross-encoder/stsb-roberta-large", + }, + ) + _register( + GreedyCrossEncoderSimilarityMatrixCalculator, + "lm_polygraph.defaults.stat_calculator_builders.default_GreedyCrossEncoderSimilarityMatrixCalculator", + { + "batch_size": 10, + "cross_encoder_name": "cross-encoder/stsb-roberta-large", + }, + ) + _register( + GreedyAlternativesNLICalculator, + "lm_polygraph.defaults.stat_calculator_builders.default_GreedyAlternativesNLICalculator", + {"nli_model": nli_model_cfg}, + ) if model_type == "Blackbox": _register(BlackboxGreedyTextsCalculator) _register(BlackboxSamplingGenerationCalculator) - if blackbox_supports_logprobs: + if model.supports_logprobs: # For blackbox models that support logprobs (like OpenAI models) _register(EntropyCalculator) - _register( - GreedyAlternativesNLICalculator, - "lm_polygraph.defaults.stat_calculator_builders.default_GreedyAlternativesNLICalculator", - {"nli_model": nli_model_cfg}, - ) - elif model_type == "Whitebox": _register( GreedyProbsCalculator, @@ -107,27 +131,6 @@ def _register( _register(SamplingPromptCalculator) _register(ClaimPromptCalculator) _register(AttentionElicitingPromptCalculator) - _register( - CrossEncoderSimilarityMatrixCalculator, - "lm_polygraph.defaults.stat_calculator_builders.default_CrossEncoderSimilarityMatrixCalculator", - { - "batch_size": deberta_batch_size, - "cross_encoder_name": "cross-encoder/stsb-roberta-large", - }, - ) - _register( - GreedyCrossEncoderSimilarityMatrixCalculator, - "lm_polygraph.defaults.stat_calculator_builders.default_GreedyCrossEncoderSimilarityMatrixCalculator", - { - "batch_size": 10, - "cross_encoder_name": "cross-encoder/stsb-roberta-large", - }, - ) - _register( - GreedyAlternativesNLICalculator, - "lm_polygraph.defaults.stat_calculator_builders.default_GreedyAlternativesNLICalculator", - {"nli_model": nli_model_cfg}, - ) _register( GreedyAlternativesFactPrefNLICalculator, "lm_polygraph.defaults.stat_calculator_builders.default_GreedyAlternativesFactPrefNLICalculator", @@ -148,6 +151,14 @@ def _register( }, ) _register(AttentionForwardPassCalculator) + _register( + TokenCrossEncoderSimilarityMatrixCalculator, + "lm_polygraph.defaults.stat_calculator_builders.default_TokenCrossEncoderSimilarityMatrixCalculator", + { + "batch_size": deberta_batch_size, + "cross_encoder_name": "cross-encoder/stsb-roberta-large", + }, + ) elif model_type == "VisualLM": _register( GreedyProbsVisualCalculator, @@ -166,19 +177,6 @@ def _register( _register(PromptVisualCalculator) _register(SamplingPromptVisualCalculator) _register(ClaimPromptVisualCalculator) - _register( - CrossEncoderSimilarityMatrixVisualCalculator, - "lm_polygraph.defaults.stat_calculator_builders.default_CrossEncoderSimilarityMatrixVisualCalculator", - { - "batch_size": deberta_batch_size, - "cross_encoder_name": "cross-encoder/stsb-roberta-large", - }, - ) - _register( - GreedyAlternativesNLICalculator, - "lm_polygraph.defaults.stat_calculator_builders.default_GreedyAlternativesNLICalculator", - {"nli_model": nli_model_cfg}, - ) _register( GreedyAlternativesFactPrefNLICalculator, "lm_polygraph.defaults.stat_calculator_builders.default_GreedyAlternativesFactPrefNLICalculator", @@ -193,6 +191,14 @@ def _register( "language": language, }, ) + _register( + TokenCrossEncoderSimilarityMatrixCalculator, + "lm_polygraph.defaults.stat_calculator_builders.default_TokenCrossEncoderSimilarityMatrixCalculator", + { + "batch_size": deberta_batch_size, + "cross_encoder_name": "cross-encoder/stsb-roberta-large", + }, + ) else: raise NotImplementedError(f"Unknown model type: {model_type}") diff --git a/src/lm_polygraph/defaults/stat_calculator_builders/default_SequenceCrossEncoderSimilarityMatrixCalculator.py b/src/lm_polygraph/defaults/stat_calculator_builders/default_SequenceCrossEncoderSimilarityMatrixCalculator.py new file mode 100644 index 000000000..b0b66641f --- /dev/null +++ b/src/lm_polygraph/defaults/stat_calculator_builders/default_SequenceCrossEncoderSimilarityMatrixCalculator.py @@ -0,0 +1,9 @@ +from lm_polygraph.stat_calculators.cross_encoder_similarity import ( + SequenceCrossEncoderSimilarityMatrixCalculator, +) + + +def load_stat_calculator(config, builder): + return SequenceCrossEncoderSimilarityMatrixCalculator( + config.batch_size, config.cross_encoder_name + ) diff --git a/src/lm_polygraph/defaults/stat_calculator_builders/default_CrossEncoderSimilarityMatrixCalculator.py b/src/lm_polygraph/defaults/stat_calculator_builders/default_TokenCrossEncoderSimilarityMatrixCalculator.py similarity index 62% rename from src/lm_polygraph/defaults/stat_calculator_builders/default_CrossEncoderSimilarityMatrixCalculator.py rename to src/lm_polygraph/defaults/stat_calculator_builders/default_TokenCrossEncoderSimilarityMatrixCalculator.py index 706b22474..daa471355 100644 --- a/src/lm_polygraph/defaults/stat_calculator_builders/default_CrossEncoderSimilarityMatrixCalculator.py +++ b/src/lm_polygraph/defaults/stat_calculator_builders/default_TokenCrossEncoderSimilarityMatrixCalculator.py @@ -1,9 +1,9 @@ from lm_polygraph.stat_calculators.cross_encoder_similarity import ( - CrossEncoderSimilarityMatrixCalculator, + TokenCrossEncoderSimilarityMatrixCalculator, ) def load_stat_calculator(config, builder): - return CrossEncoderSimilarityMatrixCalculator( + return TokenCrossEncoderSimilarityMatrixCalculator( config.batch_size, config.cross_encoder_name ) diff --git a/src/lm_polygraph/estimators/__init__.py b/src/lm_polygraph/estimators/__init__.py index e107c140b..a043fb0ed 100644 --- a/src/lm_polygraph/estimators/__init__.py +++ b/src/lm_polygraph/estimators/__init__.py @@ -77,6 +77,7 @@ from .verbalized_2s import Verbalized2S from .linguistic_1s import Linguistic1S from .claim.random_baseline import RandomBaselineClaim +from .random_baseline import RandomBaseline from .label_prob import LabelProb from .p_true_empirical import PTrueEmpirical from .focus import Focus diff --git a/src/lm_polygraph/estimators/claim/p_true.py b/src/lm_polygraph/estimators/claim/p_true.py index 9e73d0ca6..8dbaddc9b 100644 --- a/src/lm_polygraph/estimators/claim/p_true.py +++ b/src/lm_polygraph/estimators/claim/p_true.py @@ -7,7 +7,7 @@ class PTrueClaim(Estimator): """ Estimates the claim-level uncertainty of a language model following the method of "P(True)" as provided in the paper https://arxiv.org/abs/2207.05221. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). On input question q and model generation a, the method uses the following prompt: Question: {q} diff --git a/src/lm_polygraph/estimators/claim/pointwise_mutual_information.py b/src/lm_polygraph/estimators/claim/pointwise_mutual_information.py index 94153688c..5701f516e 100644 --- a/src/lm_polygraph/estimators/claim/pointwise_mutual_information.py +++ b/src/lm_polygraph/estimators/claim/pointwise_mutual_information.py @@ -9,7 +9,7 @@ class PointwiseMutualInformationClaim(Estimator): """ Estimates the claim-level uncertainty of a language model using Pointwise Mutual Information. The sequence-level estimation is calculated as product of token-level PMI estimations. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). """ def __init__(self): diff --git a/src/lm_polygraph/estimators/conditional_pointwise_mutual_information.py b/src/lm_polygraph/estimators/conditional_pointwise_mutual_information.py index 7dc23b215..542a15269 100644 --- a/src/lm_polygraph/estimators/conditional_pointwise_mutual_information.py +++ b/src/lm_polygraph/estimators/conditional_pointwise_mutual_information.py @@ -10,7 +10,7 @@ class MeanConditionalPointwiseMutualInformation(Estimator): Estimates the sequence-level uncertainty of a language model following the method of Conditional Pointwise Mutual Information (CPMI) as provided in the paper https://arxiv.org/abs/2210.13210. The sequence-level estimation is calculated as average token-level CPMI estimations. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). """ def __init__(self, tau: float = 0.0656, lambd: float = 3.599): @@ -62,7 +62,7 @@ class ConditionalPointwiseMutualInformation(Estimator): """ Estimates the token-level uncertainty of a language model following the method of Conditional Pointwise Mutual Information (CPMI) as provided in the paper https://arxiv.org/abs/2210.13210. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). """ def __init__(self, tau: float = 0.0656, lambd: float = 3.599): diff --git a/src/lm_polygraph/estimators/deg_mat.py b/src/lm_polygraph/estimators/deg_mat.py index c64ec69b5..3a4423ddf 100644 --- a/src/lm_polygraph/estimators/deg_mat.py +++ b/src/lm_polygraph/estimators/deg_mat.py @@ -12,7 +12,7 @@ class DegMat(Estimator): Estimates the sequence-level uncertainty of a language model following the method of "The Degree Matrix" as provided in the paper https://arxiv.org/abs/2305.19187. Works with both whitebox and blackbox models (initialized using - lm_polygraph.utils.model.BlackboxModel/WhiteboxModel). + lm_polygraph.model_adapters.blackbox_model.BlackboxModel/WhiteboxModel). Elements on diagonal of matrix D are sums of similarities between the particular number (position in matrix) and other answers. Thus, it is an average pairwise distance diff --git a/src/lm_polygraph/estimators/eccentricity.py b/src/lm_polygraph/estimators/eccentricity.py index fdb91cb82..537fff7d6 100644 --- a/src/lm_polygraph/estimators/eccentricity.py +++ b/src/lm_polygraph/estimators/eccentricity.py @@ -13,7 +13,7 @@ class Eccentricity(Estimator): Estimates the sequence-level uncertainty of a language model following the method of "Eccentricity" as provided in the paper https://arxiv.org/abs/2305.19187. Works with both whitebox and blackbox models (initialized using - lm_polygraph.utils.model.BlackboxModel/WhiteboxModel). + lm_polygraph.model_adapters.blackbox_model.BlackboxModel/WhiteboxModel). Method calculates a frobenious (euclidian) norm between all eigenvectors that are informative embeddings of graph Laplacian (lower norm -> closer embeddings -> higher eigenvectors -> greater uncertainty). diff --git a/src/lm_polygraph/estimators/eig_val_laplacian.py b/src/lm_polygraph/estimators/eig_val_laplacian.py index 1b8422882..aad839465 100644 --- a/src/lm_polygraph/estimators/eig_val_laplacian.py +++ b/src/lm_polygraph/estimators/eig_val_laplacian.py @@ -13,7 +13,7 @@ class EigValLaplacian(Estimator): Estimates the sequence-level uncertainty of a language model following the method of "Sum of Eigenvalues of the Graph Laplacian" as provided in the paper https://arxiv.org/abs/2305.19187. Works with both whitebox and blackbox models (initialized using - lm_polygraph.utils.model.BlackboxModel/WhiteboxModel). + lm_polygraph.model_adapters.blackbox_model.BlackboxModel/WhiteboxModel). A continuous analogue to the number of semantic sets (higher values means greater uncertainty). """ diff --git a/src/lm_polygraph/estimators/eigenscore.py b/src/lm_polygraph/estimators/eigenscore.py index c8d09f051..d59e47d87 100644 --- a/src/lm_polygraph/estimators/eigenscore.py +++ b/src/lm_polygraph/estimators/eigenscore.py @@ -9,7 +9,7 @@ class EigenScore(Estimator): """ Estimates the sequence-level uncertainty of a language model following the method of "EigenScore" as provided in the paper https://openreview.net/forum?id=Zj12nzlQbz. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). Uses embeddings for the last generated token from the middle layer of the model. """ diff --git a/src/lm_polygraph/estimators/fisher_rao.py b/src/lm_polygraph/estimators/fisher_rao.py index 84e2ac2f9..277d220d3 100644 --- a/src/lm_polygraph/estimators/fisher_rao.py +++ b/src/lm_polygraph/estimators/fisher_rao.py @@ -10,7 +10,7 @@ class FisherRao(Estimator): """ Estimates the sequence-level uncertainty of a language model following the method of "FisherRao" as provided in the paper https://arxiv.org/pdf/2212.09171.pdf. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). This method calculates the generation Fisher-Rao distance between probability distribution for each token and uniform distribution. Code adapted from https://github.com/icannos/Todd/blob/master/Todd/itscorers.py diff --git a/src/lm_polygraph/estimators/kernel_language_entropy.py b/src/lm_polygraph/estimators/kernel_language_entropy.py index 4072a469c..23d6190b4 100644 --- a/src/lm_polygraph/estimators/kernel_language_entropy.py +++ b/src/lm_polygraph/estimators/kernel_language_entropy.py @@ -55,7 +55,7 @@ class KernelLanguageEntropy(Estimator): Estimates the sequence-level uncertainty of a language model following the method of "Kernel Language Entropy" as provided in the paper https://arxiv.org/pdf/2405.20003 Works with both whitebox and blackbox models (initialized using - lm_polygraph.utils.model.BlackboxModel/WhiteboxModel). + lm_polygraph.model_adapters.blackbox_model.BlackboxModel/WhiteboxModel). This method calculates KLE(Kheat) = VNE(Kheat), where VNE is von Neumann entropy and Kheat is a heat kernel of a semantic graph over language model's outputs. diff --git a/src/lm_polygraph/estimators/lexical_similarity.py b/src/lm_polygraph/estimators/lexical_similarity.py index 43a9dab22..86669c713 100644 --- a/src/lm_polygraph/estimators/lexical_similarity.py +++ b/src/lm_polygraph/estimators/lexical_similarity.py @@ -17,7 +17,7 @@ class LexicalSimilarity(Estimator): Estimates the sequence-level uncertainty of a language model following the method of "Lexical Similarity" as provided in the paper https://arxiv.org/abs/2302.09664. Works with both whitebox and blackbox models (initialized using - lm_polygraph.utils.model.BlackboxModel/WhiteboxModel). + lm_polygraph.model_adapters.blackbox_model.BlackboxModel/WhiteboxModel). The method calculates mean similarity between all pairs of sampled generations with minus sign. The number of samples is controlled by lm_polygraph.stat_calculators.sample.SamplingGenerationCalculator diff --git a/src/lm_polygraph/estimators/max_probability.py b/src/lm_polygraph/estimators/max_probability.py index 9c3e90615..453b36d29 100644 --- a/src/lm_polygraph/estimators/max_probability.py +++ b/src/lm_polygraph/estimators/max_probability.py @@ -10,7 +10,7 @@ class MaximumSequenceProbability(Estimator): Estimates the sequence-level uncertainty of a language model by calculating the log-probability of the generation with minus sign. It is calculated as the sum of log-probabilities in each token. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). """ def __init__(self): diff --git a/src/lm_polygraph/estimators/monte_carlo_normalized_sequence_entropy.py b/src/lm_polygraph/estimators/monte_carlo_normalized_sequence_entropy.py index c7e4de5d2..7254ad6c4 100644 --- a/src/lm_polygraph/estimators/monte_carlo_normalized_sequence_entropy.py +++ b/src/lm_polygraph/estimators/monte_carlo_normalized_sequence_entropy.py @@ -10,7 +10,7 @@ def __init__(self): """ Estimates the sequence-level uncertainty of a language model following the method of "Length-normalized predictive entropy" as provided in the paper https://arxiv.org/abs/2302.09664. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). This method calculates the generation entropy estimations using Monte-Carlo estimation with length normalization. The number of samples is controlled by lm_polygraph.stat_calculators.sample.SamplingGenerationCalculator diff --git a/src/lm_polygraph/estimators/monte_carlo_sequence_entropy.py b/src/lm_polygraph/estimators/monte_carlo_sequence_entropy.py index 08de82854..7c4780a34 100644 --- a/src/lm_polygraph/estimators/monte_carlo_sequence_entropy.py +++ b/src/lm_polygraph/estimators/monte_carlo_sequence_entropy.py @@ -10,7 +10,7 @@ def __init__(self): """ Estimates the sequence-level uncertainty of a language model following the method of "Predictive entropy" as provided in the paper https://arxiv.org/abs/2302.09664. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). This method calculates the generation entropy estimations using Monte-Carlo estimation. The number of samples is controlled by lm_polygraph.stat_calculators.sample.SamplingGenerationCalculator diff --git a/src/lm_polygraph/estimators/num_sem_sets.py b/src/lm_polygraph/estimators/num_sem_sets.py index 5d8fc6068..2fb1da88e 100644 --- a/src/lm_polygraph/estimators/num_sem_sets.py +++ b/src/lm_polygraph/estimators/num_sem_sets.py @@ -14,7 +14,7 @@ class NumSemSets(Estimator): Estimates the sequence-level uncertainty of a language model following the method of "Number of Semantic Sets" as provided in the paper https://arxiv.org/abs/2305.19187. Works with both whitebox and blackbox models (initialized using - lm_polygraph.utils.model.BlackboxModel/WhiteboxModel). + lm_polygraph.model_adapters.blackbox_model.BlackboxModel/WhiteboxModel). """ def __init__( diff --git a/src/lm_polygraph/estimators/p_true.py b/src/lm_polygraph/estimators/p_true.py index 2c96eb82f..7c7621a36 100644 --- a/src/lm_polygraph/estimators/p_true.py +++ b/src/lm_polygraph/estimators/p_true.py @@ -9,7 +9,7 @@ class PTrue(Estimator): """ Estimates the sequence-level uncertainty of a language model following the method of "P(True)" as provided in the paper https://arxiv.org/abs/2207.05221. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). On input question q and model generation a, the method uses the following prompt: Question: {q} diff --git a/src/lm_polygraph/estimators/p_true_sampling.py b/src/lm_polygraph/estimators/p_true_sampling.py index eda459758..0158018f8 100644 --- a/src/lm_polygraph/estimators/p_true_sampling.py +++ b/src/lm_polygraph/estimators/p_true_sampling.py @@ -9,7 +9,7 @@ class PTrueSampling(Estimator): """ Estimates the sequence-level uncertainty of a language model following the method of "P(True)" as provided in the paper https://arxiv.org/abs/2207.05221 and model samples. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). On input question `q`, model generation `a` and several text samples `s`, the method uses the following prompt: diff --git a/src/lm_polygraph/estimators/pointwise_mutual_information.py b/src/lm_polygraph/estimators/pointwise_mutual_information.py index 0250e717d..5573155d5 100644 --- a/src/lm_polygraph/estimators/pointwise_mutual_information.py +++ b/src/lm_polygraph/estimators/pointwise_mutual_information.py @@ -9,7 +9,7 @@ class MeanPointwiseMutualInformation(Estimator): """ Estimates the sequence-level uncertainty of a language model using Pointwise Mutual Information. The sequence-level estimation is calculated as average token-level PMI estimations. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). """ def __init__(self): @@ -45,7 +45,7 @@ def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: class PointwiseMutualInformation(Estimator): """ Estimates the token-level uncertainty of a language model using Pointwise Mutual Information. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). """ def __init__(self): diff --git a/src/lm_polygraph/estimators/random_baseline.py b/src/lm_polygraph/estimators/random_baseline.py new file mode 100644 index 000000000..4f348f766 --- /dev/null +++ b/src/lm_polygraph/estimators/random_baseline.py @@ -0,0 +1,21 @@ +import numpy as np + +from typing import Dict + +from .estimator import Estimator + + +class RandomBaseline(Estimator): + """ + Sequence-level random baseline that returns scores sampled from Uniform(0, 1). + """ + + def __init__(self): + super().__init__([], "sequence") + + def __str__(self): + return "RandomBaseline" + + def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: + batch_size = len(stats["input_texts"]) + return np.random.uniform(low=0.0, high=1.0, size=batch_size) diff --git a/src/lm_polygraph/estimators/renyi_neg.py b/src/lm_polygraph/estimators/renyi_neg.py index 27bbaed6f..943442f40 100644 --- a/src/lm_polygraph/estimators/renyi_neg.py +++ b/src/lm_polygraph/estimators/renyi_neg.py @@ -10,7 +10,7 @@ class RenyiNeg(Estimator): """ Estimates the sequence-level uncertainty of a language model following the method of "RenyiNeg" as provided in the paper https://arxiv.org/pdf/2212.09171.pdf. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). This method calculates the generation Rényi divergence between probability distribution for each token and uniform distribution. Code adapted from https://github.com/icannos/Todd/blob/master/Todd/itscorers.py diff --git a/src/lm_polygraph/estimators/sar.py b/src/lm_polygraph/estimators/sar.py index 2aed3fa76..344956a15 100644 --- a/src/lm_polygraph/estimators/sar.py +++ b/src/lm_polygraph/estimators/sar.py @@ -9,7 +9,7 @@ class SAR(Estimator): """ Estimates the sequence-level uncertainty of a language model following the method of "Token SAR" as provided in the paper https://arxiv.org/abs/2307.01379. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). This method calculates the sum of corrected probability using tokenSAR of the generated text and text relevance relative to all other generations. diff --git a/src/lm_polygraph/estimators/semantic_entropy.py b/src/lm_polygraph/estimators/semantic_entropy.py index 1d0b27f14..183d38ca3 100644 --- a/src/lm_polygraph/estimators/semantic_entropy.py +++ b/src/lm_polygraph/estimators/semantic_entropy.py @@ -9,7 +9,7 @@ class SemanticEntropy(Estimator): """ Estimates the sequence-level uncertainty of a language model following the method of "Semantic entropy" as provided in the paper https://arxiv.org/abs/2302.09664. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). This method calculates the generation entropy estimations merged by semantic classes using Monte-Carlo. The number of samples is controlled by lm_polygraph.stat_calculators.sample.SamplingGenerationCalculator diff --git a/src/lm_polygraph/estimators/sentence_sar.py b/src/lm_polygraph/estimators/sentence_sar.py index 44f762afe..5ba2659b1 100644 --- a/src/lm_polygraph/estimators/sentence_sar.py +++ b/src/lm_polygraph/estimators/sentence_sar.py @@ -9,7 +9,7 @@ class SentenceSAR(Estimator): """ Estimates the sequence-level uncertainty of a language model following the method of "Sentence SAR" as provided in the paper https://arxiv.org/abs/2307.01379. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). This method calculates the sum of the probability of the generated text and text relevance relative to all other generations. """ diff --git a/src/lm_polygraph/estimators/token_entropy.py b/src/lm_polygraph/estimators/token_entropy.py index fc87cc77c..374981f7e 100644 --- a/src/lm_polygraph/estimators/token_entropy.py +++ b/src/lm_polygraph/estimators/token_entropy.py @@ -9,7 +9,7 @@ class MeanTokenEntropy(Estimator): """ Estimates the sequence-level uncertainty of a language model by calculating the mean entropy among all tokens in the generation. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). """ def __init__(self): @@ -37,7 +37,7 @@ class TokenEntropy(Estimator): """ Estimates the token-level uncertainty of a language model by calculating the entropy for each token in the generation. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). """ def __init__(self): diff --git a/src/lm_polygraph/estimators/token_sar.py b/src/lm_polygraph/estimators/token_sar.py index 078e16458..21f1c256d 100644 --- a/src/lm_polygraph/estimators/token_sar.py +++ b/src/lm_polygraph/estimators/token_sar.py @@ -25,7 +25,7 @@ class TokenSAR(Estimator): """ Estimates the sequence-level uncertainty of a language model following the method of "Token SAR" as provided in the paper https://arxiv.org/abs/2307.01379. - Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Works only with whitebox models (initialized using lm_polygraph.model_adapters.whitebox_model.WhiteboxModel). This method calculates the weighted sum of log_likelihoods with weights computed using token relevance. """ diff --git a/src/lm_polygraph/model_adapters/__init__.py b/src/lm_polygraph/model_adapters/__init__.py index b5be0fef9..0876241d0 100644 --- a/src/lm_polygraph/model_adapters/__init__.py +++ b/src/lm_polygraph/model_adapters/__init__.py @@ -1,5 +1,9 @@ from .whitebox_model_basic import WhiteboxModelBasic from .visual_whitebox_model import VisualWhiteboxModel from .whitebox_model_vllm import WhiteboxModelvLLM -from lm_polygraph.utils.model import WhiteboxModel -from lm_polygraph.utils.model import BlackboxModel +from .whitebox_model import WhiteboxModel +from .blackbox_model import BlackboxModel + +from . import openai_adapter +from . import togetherai_adapter +from . import huggingface_adapter diff --git a/src/lm_polygraph/model_adapters/api_provider_adapter.py b/src/lm_polygraph/model_adapters/api_provider_adapter.py new file mode 100644 index 000000000..b32f1398a --- /dev/null +++ b/src/lm_polygraph/model_adapters/api_provider_adapter.py @@ -0,0 +1,140 @@ +""" +API Provider Adapter System + +This module implements a flexible adapter pattern for handling different API providers +with varying parameter formats and response structures. +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import List, Dict, Optional, Any +import logging + +log = logging.getLogger("lm_polygraph") + +# The registry to hold all available adapters +ADAPTER_REGISTRY = {} + + +def register_adapter(api_provider_name: str): + """A decorator to register an adapter class.""" + + def decorator(cls): + ADAPTER_REGISTRY[api_provider_name] = cls + return cls + + return decorator + + +@dataclass +class StandardizedResponse: + """A unified data structure for API responses.""" + + text: str + tokens: Optional[List[str]] = None + logprobs: Optional[List[float]] = None + top_logprobs: Optional[List[List[float]]] = None + alternative_tokens: Optional[List[List[str]]] = None + finish_reason: Optional[str] = None + raw_response: Optional[Dict[str, Any]] = None + + +class APIProviderAdapter(ABC): + """ + Abstract base class for API Provider Adapters. + + Each adapter is responsible for: + 1. Translating request parameters to provider-specific format + 2. Parsing provider-specific responses to standardized format + """ + + @abstractmethod + def adapt_request(self, params: dict) -> dict: + """ + Adapts the generation parameters for the specific API provider. + + Args: + params: Dictionary of generation parameters + + Returns: + Modified parameters dictionary suitable for the provider + """ + pass + + @abstractmethod + def parse_response(self, response: dict) -> StandardizedResponse: + """ + Parses the raw API response into the StandardizedResponse format. + + Args: + response: Raw API response dictionary + + Returns: + StandardizedResponse object + """ + pass + + @abstractmethod + def generate_texts( + self, + model, + input_texts: List[Any], + args: Dict[str, Any], + ) -> List[Any]: + """Execute provider-specific inference. + + Args: + model: The calling BlackboxModel instance for state updates. + input_texts: List of prompts/messages to generate for. + args: Provider-specific generation parameters. + + Returns: + List of generated outputs matching the provider's format. + """ + pass + + def supports_logprobs(self, model_path: str = None) -> bool: + """ + Check if the provider/model supports logprobs. + + Args: + model_path: Optional model identifier for model-specific checks + + Returns: + True if logprobs are supported, False otherwise + """ + return False + + def validate_parameter_ranges(self, params: dict) -> dict: + """ + Validate and clamp parameters to provider-specific ranges. + + Args: + params: Parameters dictionary + + Returns: + Validated parameters dictionary + """ + return params + + +def get_adapter(api_provider_name: str) -> APIProviderAdapter: + """ + Factory function to get an adapter instance from the registry. + + Args: + api_provider_name: Name of the API provider + + Returns: + APIProviderAdapter instance + + Defaults to 'openai' if the requested provider is not found. + """ + if api_provider_name not in ADAPTER_REGISTRY: + raise ValueError( + f"No adapter found for provider '{api_provider_name}'. Available providers: {list(ADAPTER_REGISTRY.keys())}" + ) + + adapter_class = ADAPTER_REGISTRY[api_provider_name] + + return adapter_class() diff --git a/src/lm_polygraph/model_adapters/blackbox_model.py b/src/lm_polygraph/model_adapters/blackbox_model.py new file mode 100644 index 000000000..a0f9ba34a --- /dev/null +++ b/src/lm_polygraph/model_adapters/blackbox_model.py @@ -0,0 +1,142 @@ +from dataclasses import asdict +from typing import List, Optional + +from lm_polygraph.utils.model import Model +from lm_polygraph.utils.generation_parameters import GenerationParameters +from lm_polygraph.model_adapters.api_provider_adapter import get_adapter + + +class BlackboxModel(Model): + """ + Black-box model class. Have no access to model scores and logits. + Currently implemented blackbox models: OpenAI models, Huggingface models. + + Examples: + + ```python + >>> from lm_polygraph import BlackboxModel + >>> model = BlackboxModel( + ... model_path='gpt-3.5-turbo', + ... api_provider_name='openai' + ... ) + ``` + + ```python + >>> from lm_polygraph import BlackboxModel + >>> model = BlackboxModel( + ... model_path='google/t5-large-ssm-nqo', + ... api_provider_name='huggingface' + ... ) + ``` + """ + + def __init__( + self, + model_path: str = None, + generation_parameters: GenerationParameters = GenerationParameters(), + api_provider_name: str = "openai", + tokenizer: object = None, + supports_logprobs: Optional[bool] = None, + ): + """ + Parameters: + model_path (Optional[str]): Unique model path or identifier understood by the provider. + generation_parameters (GenerationParameters): parameters to use in model generation. Default: default parameters. + supports_logprobs (Optional[bool]): Override adapter capability detection for logprobs. + If None, adapter default behavior is used. + """ + super().__init__(model_path, "Blackbox") + self.generation_parameters = generation_parameters + self.api_provider_name = api_provider_name + self._supports_logprobs_override = supports_logprobs + + # Initialize the adapter for this provider + self.adapter = get_adapter(self.api_provider_name) + self.tokenizer = tokenizer + + @property + def supports_logprobs(self) -> bool: + """Expose adapter-defined logprob support.""" + if self._supports_logprobs_override is not None: + return self._supports_logprobs_override + return self.adapter.supports_logprobs(self.model_path) + + def _validate_args(self, args): + """ + Validates and adapts arguments for BlackboxModel generation using the configured adapter. + + Parameters: + args (dict): The arguments to validate. + + Returns: + dict: Validated and adapted arguments. + """ + # Use the adapter to validate parameter ranges first + validated_args = self.adapter.validate_parameter_ranges(args) + + # Then adapt the request format for the specific provider + return self.adapter.adapt_request(validated_args) + + def prepare_input(self, prompt: List[str]) -> List[dict]: + if isinstance(prompt, str): + messages = [{"role": "user", "content": prompt}] + elif isinstance(prompt, list) and all( + isinstance(item, dict) for item in prompt + ): + messages = prompt + else: + raise ValueError( + "Invalid prompt format. Must be either a string or a list of dictionaries." + ) + + return messages + + def generate_texts(self, input_texts: List[str], **args) -> List[str]: + """ + Generates a list of model answers using input texts batch. + + Parameters: + input_texts (List[str]): input texts batch. + Return: + List[str]: corresponding model generations. Have the same length as `input_texts`. + """ + # Apply default parameters first, then override with provided args + default_params = asdict(self.generation_parameters) + default_params.update(args) + args = self._validate_args(default_params) + + # Check if we're trying to access features that require logprobs support + requires_logprobs = args.get("output_scores", False) + + # Use adapter to check logprobs support (considering model-specific rules) + if requires_logprobs and not self.supports_logprobs: + raise Exception( + f"Cannot access logits for blackbox model with provider '{self.api_provider_name}'" + ) + + return self.adapter.generate_texts(self, input_texts, args) + + def generate(self, **args): + """ + Generates a single model answer using greedy decoding. + Parameters: + args: input arguments for generation. + Returns: + List[str]: corresponding model generations. Have the same length as `input_texts`. + """ + args["temperature"] = 0.0 + args["n"] = 1 + + output = self.generate_texts(**args) + + # For each input in batch model returns a list with one generation + # we take the first generation from each list + output = [out[0] for out in output] + + return output + + def __call__(self, **args): + """ + Not implemented for blackbox models. + """ + raise Exception("Cannot call blackbox model directly, use generate or generate_texts methods.") diff --git a/src/lm_polygraph/model_adapters/huggingface_adapter.py b/src/lm_polygraph/model_adapters/huggingface_adapter.py new file mode 100644 index 000000000..302a29eae --- /dev/null +++ b/src/lm_polygraph/model_adapters/huggingface_adapter.py @@ -0,0 +1,141 @@ +"""Hugging Face Inference API adapter.""" + +import logging +from typing import Any, List + +from huggingface_hub import InferenceClient + +from .api_provider_adapter import ( + APIProviderAdapter, + StandardizedResponse, + register_adapter, +) + +log = logging.getLogger("lm_polygraph") + + +@register_adapter("huggingface") +class HuggingFaceAdapter(APIProviderAdapter): + """Minimal adapter for Hugging Face text generation endpoints.""" + + def adapt_request(self, params: dict) -> dict: + """ + Adapts parameters for HF API format. + """ + args_copy = params.copy() + + # BlackboxModel specific validation - remove unsupported parameters + for delete_key in [ + "do_sample", + "min_length", + "top_k", + "repetition_penalty", + "min_new_tokens", + "num_beams", + "allow_newlines", + ]: + args_copy.pop(delete_key, None) + + key_mapping = { + "num_return_sequences": "n", + "max_length": "max_tokens", + "max_new_tokens": "max_tokens", + "stop_strings": "stop", + } + for key, replace_key in key_mapping.items(): + if key in args_copy: + args_copy[replace_key] = args_copy[key] + args_copy.pop(key) + + return args_copy + + def parse_response(self, response: dict) -> StandardizedResponse: + """Parse the response into the standardized structure.""" + text = response.message.content + + # Extract logprobs if available + logprobs = None + tokens = None + top_logprobs = None + alternative_tokens = None + if hasattr(response, "logprobs") and response.logprobs: + logprobs_data = response.logprobs + if hasattr(logprobs_data, "content") and logprobs_data.content: + # Extract tokens from logprobs content + for item in logprobs_data.content: + if hasattr(item, "token"): + tokens = tokens or [] + tokens.append(item.token) + + if hasattr(item, "logprob"): + logprobs = logprobs or [] + logprobs.append(item.logprob) + + if hasattr(item, "top_logprobs"): + top_logprobs = top_logprobs or [] + top_logprobs.append( + [item.logprob for pair in item.top_logprobs] + ) + + alternative_tokens = alternative_tokens or [] + alternative_tokens.append( + [pair.token for pair in item.top_logprobs] + ) + + return StandardizedResponse( + text=text, + tokens=tokens, + logprobs=logprobs, + top_logprobs=top_logprobs, + alternative_tokens=alternative_tokens, + finish_reason=response.finish_reason, + raw_response=response, + ) + + def supports_logprobs(self, model_path: str = None) -> bool: + return True + + def generate_texts( + self, + model, + input_texts: List[Any], + args: dict, + ) -> List[Any]: + if model.model_path is None: + raise ValueError( + "model_path must be specified for Huggingface API inference." + ) + + client = InferenceClient(model=model.model_path) + + return_logprobs = args.pop("output_scores", False) + logprobs_args = {} + + if return_logprobs and model.supports_logprobs: + logprobs_args["logprobs"] = True + logprobs_args["top_logprobs"] = args.pop("top_logprobs", 5) + + parsed_responses = [] + for prompt in input_texts: + messages = model.prepare_input(prompt) + + retries = 0 + while True: + try: + response = client.chat_completion( + messages=messages, + **args, + **logprobs_args, + ) + break + except Exception as e: + if retries > 4: + raise Exception from e + retries += 1 + continue + + parsed_responses.append( + [self.parse_response(resp) for resp in response.choices] + ) + + return parsed_responses diff --git a/src/lm_polygraph/model_adapters/openai_adapter.py b/src/lm_polygraph/model_adapters/openai_adapter.py new file mode 100644 index 000000000..8988512ab --- /dev/null +++ b/src/lm_polygraph/model_adapters/openai_adapter.py @@ -0,0 +1,211 @@ +"""OpenAI-compatible API Provider Adapter implementations.""" + +import logging +from typing import Any, List + +import openai + +from .api_provider_adapter import ( + APIProviderAdapter, + StandardizedResponse, + register_adapter, +) + +log = logging.getLogger("lm_polygraph") + + +class OpenAIChatCompletionMixin: + """Reusable chat completion inference flow for OpenAI-compatible providers.""" + + def _create_client(self, model): + raise NotImplementedError + + def generate_texts( + self, + model, + input_texts: List[Any], + args: dict, + ) -> List[Any]: + openai_api = self._create_client(model) + + return_logprobs = args.pop("output_scores", False) + logprobs_args = {} + + if return_logprobs: + logprobs_args["logprobs"] = True + logprobs_args["top_logprobs"] = args.pop("top_logprobs", 5) + + parsed_responses = [] + for prompt in input_texts: + messages = model.prepare_input(prompt) + + retries = 0 + while True: + try: + response = openai_api.chat.completions.create( + model=model.model_path, + messages=messages, + **args, + **logprobs_args, + ) + break + except Exception as e: + if retries > 4: + raise Exception from e + retries += 1 + continue + + parsed_responses.append( + [self.parse_response(resp) for resp in response.choices] + ) + + return parsed_responses + + +@register_adapter("openai") +class OpenAIAdapter(OpenAIChatCompletionMixin, APIProviderAdapter): + """ + Adapter for OpenAI API. + + This adapter replicates the exact behavior of the original BlackboxModel + _validate_args method and response parsing to ensure backward compatibility. + """ + + def adapt_request(self, params: dict) -> dict: + """ + Adapts parameters for OpenAI API format. + """ + args_copy = params.copy() + + # BlackboxModel specific validation - remove unsupported parameters + for delete_key in [ + "do_sample", + "min_length", + "top_k", + "repetition_penalty", + "min_new_tokens", + "num_beams", + "stop_strings", + "allow_newlines", + ]: + args_copy.pop(delete_key, None) + + # Map HF argument names to OpenAI API argument names + key_mapping = { + "num_return_sequences": "n", + "max_length": "max_completion_tokens", + "max_new_tokens": "max_completion_tokens", + } + for key, replace_key in key_mapping.items(): + if key in args_copy: + args_copy[replace_key] = args_copy[key] + args_copy.pop(key) + + return args_copy + + def parse_response(self, response: dict) -> StandardizedResponse: + """ + Parses OpenAI API response into standardized format. + + Args: + response: Raw OpenAI API response dictionary + + Returns: + StandardizedResponse object + """ + try: + # Extract main text content + text = response.message.content + + # Extract logprobs if available + logprobs = None + tokens = None + top_logprobs = None + alternative_tokens = None + if hasattr(response, "logprobs") and response.logprobs: + logprobs_data = response.logprobs + if hasattr(logprobs_data, "content") and logprobs_data.content: + # Extract tokens from logprobs content + for item in logprobs_data.content: + if hasattr(item, "token"): + tokens = tokens or [] + tokens.append(item.token) + + if hasattr(item, "logprob"): + logprobs = logprobs or [] + logprobs.append(item.logprob) + + if hasattr(item, "top_logprobs"): + top_logprobs = top_logprobs or [] + top_logprobs.append( + [item.logprob for item in item.top_logprobs] + ) + + alternative_tokens = alternative_tokens or [] + alternative_tokens.append( + [pair.token for pair in item.top_logprobs] + ) + + return StandardizedResponse( + text=text, + tokens=tokens, + logprobs=logprobs, + top_logprobs=top_logprobs, + alternative_tokens=alternative_tokens, + finish_reason=response.finish_reason, + raw_response=response, + ) + + except (AttributeError, IndexError, TypeError) as e: + log.error(f"Error parsing OpenAI response: {e}") + log.error(f"Response structure: {response}") + raise ValueError(f"Invalid OpenAI API response format: {e}") + + def supports_logprobs(self, model_path: str = None) -> bool: + """ + Check if the OpenAI model supports logprobs. + + Args: + model_path: OpenAI model identifier + + Returns: + True (OpenAI generally supports logprobs) + """ + return True + + def validate_parameter_ranges(self, params: dict) -> dict: + """ + Validate and clamp parameters to OpenAI-specific ranges. + + Args: + params: Parameters dictionary + + Returns: + Validated parameters dictionary + """ + validated_params = params.copy() + + # OpenAI parameter ranges + parameter_ranges = { + "temperature": (0.0, 2.0), + "top_p": (0.0, 1.0), + "presence_penalty": (-2.0, 2.0), + "frequency_penalty": (-2.0, 2.0), + } + + for param, (min_val, max_val) in parameter_ranges.items(): + if param in validated_params: + value = validated_params[param] + if isinstance(value, (int, float)): + original_value = value + validated_params[param] = max(min_val, min(max_val, value)) + if original_value != validated_params[param]: + log.warning( + f"Parameter {param}={original_value} clamped to " + f"{validated_params[param]} (range: [{min_val}, {max_val}]) for OpenAI" + ) + + return validated_params + + def _create_client(self, model): + return openai.OpenAI() diff --git a/src/lm_polygraph/model_adapters/togetherai_adapter.py b/src/lm_polygraph/model_adapters/togetherai_adapter.py new file mode 100644 index 000000000..ebdc9626f --- /dev/null +++ b/src/lm_polygraph/model_adapters/togetherai_adapter.py @@ -0,0 +1,243 @@ +"""Together.ai API Provider Adapter.""" + +import logging +import os +from typing import Any, List + +from together import Together +import numpy as np + +from .api_provider_adapter import ( + APIProviderAdapter, + StandardizedResponse, + register_adapter, +) + +log = logging.getLogger("lm_polygraph") + + +class TogetherAIChatCompletionMixin: + """Reusable chat completion inference flow for TogetherAI-compatible providers.""" + + def _create_client(self, model): + raise NotImplementedError + + def generate_texts( + self, + model, + input_texts: List[Any], + args: dict, + ) -> List[Any]: + together_api = self._create_client(model) + + return_logprobs = args.pop("output_scores", False) + logprobs_args = {} + + if return_logprobs: + logprobs_args["logprobs"] = args.pop("top_logprobs", 5) + + parsed_responses = [] + for prompt in input_texts: + messages = model.prepare_input(prompt) + + retries = 0 + while True: + try: + response = together_api.chat.completions.create( + model=model.model_path, + messages=messages, + **args, + **logprobs_args, + ) + break + except Exception as e: + if retries > 4: + raise Exception from e + retries += 1 + continue + + parsed_responses.append( + [self.parse_response(resp) for resp in response.choices] + ) + + return parsed_responses + +@register_adapter("together_ai") +class TogetherAIAdapter(TogetherAIChatCompletionMixin, APIProviderAdapter): + """ + Adapter for Together.ai API. + + Together.ai provides an OpenAI-compatible API with some enhancements + like additional logprobs features and model-specific parameters. + """ + + def adapt_request(self, params: dict) -> dict: + """ + Adapts parameters for together.ai API format. + + Together.ai is highly compatible with OpenAI API format, + so most parameters can be passed through as-is. + """ + adapted_params = params.copy() + + # Remove parameters that together.ai doesn't support or handles differently + for delete_key in [ + "do_sample", + "min_length", + "min_new_tokens", + "num_beams", + "stop_strings", + "allow_newlines", + "top_k", # together.ai doesn't accept top_k through OpenAI client + "repetition_penalty", # Use frequency_penalty instead + ]: + adapted_params.pop(delete_key, None) + + # Map HF argument names to together.ai API argument names + key_mapping = { + "num_return_sequences": "n", + "max_length": "max_tokens", + "max_new_tokens": "max_tokens", + } + for key, replace_key in key_mapping.items(): + if key in adapted_params: + adapted_params[replace_key] = adapted_params[key] + adapted_params.pop(key) + + return adapted_params + + def parse_response(self, response: dict) -> StandardizedResponse: + """ + Parses together.ai API response into standardized format. + + Together.ai follows OpenAI format very closely, with some additional fields. + + Args: + response: Raw together.ai API response dictionary + + Returns: + StandardizedResponse object + """ + try: + # Extract main text content (same as OpenAI format) + text = response.message.content + + # Extract logprobs if available - transform to OpenAI format for stat calculator compatibility + logprobs = None + tokens = None + top_logprobs = None + alternative_tokens = None + if hasattr(response, "logprobs") and response.logprobs: + logprobs_data = response.logprobs + + if hasattr(logprobs_data, "content"): + # for some models, together.ai returns logprobs nested under 'content' + content_logprobs = logprobs_data.content + + logprobs = [] + tokens = [] + top_logprobs = [] + alternative_tokens = [] + for item in content_logprobs: + if "token" in item: + tokens.append(item["token"]) + if "logprob" in item: + logprobs.append(item["logprob"]) + if "top_logprobs" in item: + tl_dict = item["top_logprobs"] + alternative_tokens.append([tl_item['token'] for tl_item in tl_dict]) + top_logprobs.append([tl_item['logprob'] for tl_item in tl_dict]) + else: + if hasattr(logprobs_data, "token_ids"): + tokens = logprobs_data.token_ids + if hasattr(logprobs_data, "token_logprobs"): + logprobs = logprobs_data.token_logprobs + if hasattr(logprobs_data, "top_logprobs"): + top_logprobs = [ + list(tl_dict.values()) for tl_dict in logprobs_data.top_logprobs + ] + alternative_tokens = [ + list(tl_dict.keys()) for tl_dict in logprobs_data.top_logprobs + ] + + max_num_lp = np.max([len(lp) for lp in top_logprobs]) + min_num_lp = np.min([len(lp) for lp in top_logprobs]) + + if max_num_lp != min_num_lp: + # clip all to min length for consistency + top_logprobs = [ + lp[:min_num_lp] for lp in top_logprobs + ] + alternative_tokens = [ + at[:min_num_lp] for at in alternative_tokens + ] + + return StandardizedResponse( + text=text, + tokens=tokens, + logprobs=logprobs, + top_logprobs=top_logprobs, + alternative_tokens=alternative_tokens, + finish_reason=response.finish_reason, + raw_response=response, + ) + + except (KeyError, IndexError, TypeError) as e: + log.error(f"Error parsing together.ai response: {e}") + log.error(f"Response structure: {response}") + raise ValueError(f"Invalid together.ai API response format: {e}") + + def supports_logprobs(self, model_path: str = None) -> bool: + """ + Check if the together.ai model supports logprobs. + + Args: + model_path: Together.ai model identifier + + Returns: + True (together.ai generally supports logprobs) + """ + return True + + def validate_parameter_ranges(self, params: dict) -> dict: + """ + Validate and clamp parameters to together.ai-specific ranges. + + Args: + params: Parameters dictionary + + Returns: + Validated parameters dictionary + """ + validated_params = params.copy() + + # Together.ai parameter ranges (same as OpenAI with some additions) + parameter_ranges = { + "temperature": (0.0, 2.0), + "top_p": (0.0, 1.0), + "top_k": (1, 200), # Together.ai specific + "presence_penalty": (-2.0, 2.0), + "frequency_penalty": (-2.0, 2.0), + "repetition_penalty": (0.0, 2.0), # Together.ai specific + } + + for param, (min_val, max_val) in parameter_ranges.items(): + if param in validated_params: + value = validated_params[param] + if isinstance(value, (int, float)): + original_value = value + validated_params[param] = max(min_val, min(max_val, value)) + if original_value != validated_params[param]: + log.warning( + f"Parameter {param}={original_value} clamped to " + f"{validated_params[param]} (range: [{min_val}, {max_val}]) for together.ai" + ) + + return validated_params + + def _create_client(self, model): + client = Together( + api_key=os.environ.get("TOGETHER_API_KEY"), + ) + + return client diff --git a/src/lm_polygraph/model_adapters/whitebox_model.py b/src/lm_polygraph/model_adapters/whitebox_model.py new file mode 100644 index 000000000..54e1a16f6 --- /dev/null +++ b/src/lm_polygraph/model_adapters/whitebox_model.py @@ -0,0 +1,293 @@ +import torch +import logging + +from dataclasses import asdict +from typing import List, Dict, Optional, Union +from transformers import ( + AutoTokenizer, + AutoModelForSeq2SeqLM, + AutoModelForCausalLM, + AutoConfig, + LogitsProcessorList, + BartForConditionalGeneration, +) + +from lm_polygraph.utils.model import Model +from lm_polygraph.utils.generation_parameters import ( + GenerationParameters, + GenerationParametersFactory, +) + +log = logging.getLogger("lm_polygraph") + + +class WhiteboxModel(Model): + """ + White-box model class. Have access to model scores and logits. Currently implemented only for Huggingface models. + + Examples: + + ```python + >>> from lm_polygraph import WhiteboxModel + >>> model = WhiteboxModel.from_pretrained( + ... "bigscience/bloomz-3b", + ... ) + ``` + """ + + def __init__( + self, + model: AutoModelForCausalLM, + tokenizer: AutoTokenizer, + model_path: str = None, + model_type: str = "CausalLM", + generation_parameters: GenerationParameters = GenerationParameters(), + instruct: bool = False, + ): + """ + Parameters: + model (AutoModelForCausalLM): HuggingFace model. + tokenizer (AutoTokenizer): HuggingFace tokenizer. + model_path (Optional[str]): Unique model path in HuggingFace. + model_type (str): Additional model specifications. + parameters (GenerationParameters): parameters to use in model generation. Default: default parameters. + """ + super().__init__(model_path, model_type) + self.model = model + self.tokenizer = tokenizer + self.generation_parameters = generation_parameters + self.instruct = instruct + + def _validate_args(self, args): + """ + Validates and adapts arguments for WhiteboxModel generation. + + Parameters: + args (dict): The arguments to validate. + + Returns: + dict: Validated and adapted arguments. + """ + args_copy = args.copy() + + # WhiteboxModel specific validation + if "presence_penalty" in args_copy and args_copy["presence_penalty"] != 0.0: + log.warning( + "Skipping requested argument presence_penalty={}".format( + args_copy["presence_penalty"] + ) + ) + + # Remove arguments that are not supported by the HF model.generate function + keys_to_remove = ["presence_penalty", "allow_newlines"] + for key in keys_to_remove: + args_copy.pop(key, None) + + return args_copy + + class _ScoresProcessor: + # Stores original token scores instead of the ones modified with generation parameters + def __init__(self): + self.scores = [] + + def __call__(self, input_ids=None, scores=None): + self.scores.append(scores.log_softmax(-1)) + return scores + + def generate(self, **args): + """ + Generates the model output with scores from batch formed by HF Tokenizer. + + Parameters: + **args: Any arguments that can be passed to model.generate function from HuggingFace. + Returns: + ModelOutput: HuggingFace generation output with scores overriden with original probabilities. + """ + default_params = asdict(self.generation_parameters) + + # add ScoresProcessor to collect original scores + processor = self._ScoresProcessor() + if "logits_processor" in args.keys(): + logits_processor = LogitsProcessorList( + [processor, args["logits_processor"]] + ) + else: + logits_processor = LogitsProcessorList([processor]) + args["logits_processor"] = logits_processor + + # update default parameters with passed arguments + default_params.update(args) + args = default_params + + if "stop_strings" in args: + args["tokenizer"] = self.tokenizer + + args = self._validate_args(args) + generation = self.model.generate(**args) + + # override generation.scores with original scores from model + generation.generation_scores = generation.scores + generation.scores = processor.scores + + return generation + + def generate_texts(self, input_texts: List[str], **args) -> List[str]: + """ + Generates a list of model answers using input texts batch. + + Parameters: + input_texts (List[str]): input texts batch. + Return: + List[str]: corresponding model generations. Have the same length as `input_texts`. + """ + # Apply default parameters first, then override with provided args + default_params = asdict(self.generation_parameters) + default_params.update(args) + args = self._validate_args(default_params) + + args["return_dict_in_generate"] = True + batch: Dict[str, torch.Tensor] = self.tokenize(input_texts) + batch = {k: v.to(self.device()) for k, v in batch.items()} + sequences = self.generate(**batch, **args).sequences.cpu() + input_len = batch["input_ids"].shape[1] + texts = [] + + decode_args = {} + if self.tokenizer.chat_template is not None: + decode_args["skip_special_tokens"] = True + + for seq in sequences: + if self.model_type == "CausalLM": + texts.append(self.tokenizer.decode(seq[input_len:], **decode_args)) + else: + texts.append(self.tokenizer.decode(seq[1:], **decode_args)) + + return texts + + def __call__(self, **args): + """ + Calls the model on the input batch. Returns the resulted scores. + """ + return self.model(**args) + + def device(self): + """ + Returns the device the model is currently loaded on. + + Returns: + str: device string. + """ + return self.model.device + + @staticmethod + def from_pretrained( + model_path: str, + generation_params: Optional[Dict] = {}, + add_bos_token: bool = True, + **kwargs, + ): + """ + Initializes the model from HuggingFace. Automatically determines model type. + + Parameters: + model_path (str): model path in HuggingFace. + generation_params (Dict): generation arguments for + lm_polygraph.utils.generation_parametersGenerationParameters + add_bos_token (bool): tokenizer argument. Default: True. + """ + log.warning( + "WhiteboxModel#from_pretrained is deprecated and will be removed in the next release. Please instantiate WhiteboxModel directly by passing an already loaded model, tokenizer and model path." + ) + + config = AutoConfig.from_pretrained( + model_path, trust_remote_code=True, **kwargs + ) + + if any(["CausalLM" in architecture for architecture in config.architectures]): + model_type = "CausalLM" + model = AutoModelForCausalLM.from_pretrained( + model_path, trust_remote_code=True, **kwargs + ) + elif any( + [ + ("Seq2SeqLM" in architecture) + or ("ConditionalGeneration" in architecture) + for architecture in config.architectures + ] + ): + model_type = "Seq2SeqLM" + model = AutoModelForSeq2SeqLM.from_pretrained(model_path, **kwargs) + if "falcon" in model_path: + model.transformer.alibi = True + elif any( + ["JAISLMHeadModel" in architecture for architecture in config.architectures] + ): + model_type = "CausalLM" + model = AutoModelForCausalLM.from_pretrained( + model_path, + trust_remote_code=True, + **kwargs, + ) + elif any( + ["BartModel" in architecture for architecture in config.architectures] + ): + model_type = "Seq2SeqLM" + model = BartForConditionalGeneration.from_pretrained(model_path, **kwargs) + else: + raise ValueError( + f"Model {model_path} is not adapted for the sequence generation task" + ) + + tokenizer = AutoTokenizer.from_pretrained( + model_path, + padding_side="left", + add_bos_token=add_bos_token, + **kwargs, + ) + + model.eval() + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + generation_params = GenerationParametersFactory.from_params( + yaml_config=generation_params, + native_config=asdict(model.config), + ) + + instance = WhiteboxModel( + model, tokenizer, model_path, model_type, generation_params + ) + + return instance + + def tokenize( + self, texts: Union[List[str], List[List[Dict[str, str]]]] + ) -> Dict[str, torch.Tensor]: + """ + Tokenizes input texts batch into a dictionary using the model tokenizer. + + Parameters: + texts (List[str]): list of input texts batch. + Returns: + dict[str, torch.Tensor]: tensors dictionary obtained by tokenizing input texts batch. + """ + # Apply chat template if tokenizer has it + add_start_symbol = True + if self.instruct: + formatted_texts = [] + for chat in texts: + if isinstance(chat, str): + chat = [{"role": "user", "content": chat}] + formatted_chat = self.tokenizer.apply_chat_template( + chat, add_generation_prompt=True, tokenize=False + ) + formatted_texts.append(formatted_chat) + texts = formatted_texts + + add_start_symbol = False + return self.tokenizer( + texts, + padding=True, + return_tensors="pt", + add_special_tokens=add_start_symbol, + ) diff --git a/src/lm_polygraph/stat_calculators/__init__.py b/src/lm_polygraph/stat_calculators/__init__.py index 871a8c49d..9e0e3c26f 100644 --- a/src/lm_polygraph/stat_calculators/__init__.py +++ b/src/lm_polygraph/stat_calculators/__init__.py @@ -36,8 +36,8 @@ from .entropy import EntropyCalculator from .sample import ( SamplingGenerationCalculator, - BlackboxSamplingGenerationCalculator, ) +from .sample_blackbox import BlackboxSamplingGenerationCalculator from .greedy_alternatives_nli import ( GreedyAlternativesNLICalculator, GreedyAlternativesFactPrefNLICalculator, @@ -53,7 +53,10 @@ GreedySemanticMatrixCalculator, ConcatGreedySemanticMatrixCalculator, ) -from .cross_encoder_similarity import CrossEncoderSimilarityMatrixCalculator +from .cross_encoder_similarity import ( + SequenceCrossEncoderSimilarityMatrixCalculator, + TokenCrossEncoderSimilarityMatrixCalculator, +) from .greedy_cross_encoder_similarity import ( GreedyCrossEncoderSimilarityMatrixCalculator, ) diff --git a/src/lm_polygraph/stat_calculators/attention_eliciting_prompt.py b/src/lm_polygraph/stat_calculators/attention_eliciting_prompt.py index cab55a273..0d267c6f9 100644 --- a/src/lm_polygraph/stat_calculators/attention_eliciting_prompt.py +++ b/src/lm_polygraph/stat_calculators/attention_eliciting_prompt.py @@ -4,7 +4,7 @@ from typing import Dict, List, Tuple from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel class BaseAttentionPromptCalculator(StatCalculator): diff --git a/src/lm_polygraph/stat_calculators/attention_forward_pass.py b/src/lm_polygraph/stat_calculators/attention_forward_pass.py index 0e430bc27..4800c0272 100644 --- a/src/lm_polygraph/stat_calculators/attention_forward_pass.py +++ b/src/lm_polygraph/stat_calculators/attention_forward_pass.py @@ -4,7 +4,7 @@ from typing import Dict, List, Tuple from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel class AttentionForwardPassCalculator(StatCalculator): diff --git a/src/lm_polygraph/stat_calculators/bart_score.py b/src/lm_polygraph/stat_calculators/bart_score.py index b60e5df48..327c1356f 100644 --- a/src/lm_polygraph/stat_calculators/bart_score.py +++ b/src/lm_polygraph/stat_calculators/bart_score.py @@ -7,7 +7,7 @@ import torch.nn as nn from transformers import BartTokenizer, BartForConditionalGeneration from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel log = logging.getLogger(__name__) diff --git a/src/lm_polygraph/stat_calculators/cross_encoder_similarity.py b/src/lm_polygraph/stat_calculators/cross_encoder_similarity.py index 893f9e1ad..420e44aa8 100644 --- a/src/lm_polygraph/stat_calculators/cross_encoder_similarity.py +++ b/src/lm_polygraph/stat_calculators/cross_encoder_similarity.py @@ -1,14 +1,16 @@ import numpy as np +import torch import itertools from typing import Dict, List, Tuple from .stat_calculator import StatCalculator from sentence_transformers import CrossEncoder -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters import * +from lm_polygraph.utils.model import Model -class CrossEncoderSimilarityMatrixCalculator(StatCalculator): +class SequenceCrossEncoderSimilarityMatrixCalculator(StatCalculator): """ Calculates the cross-encoder similarity matrix for generation samples using RoBERTa model. """ @@ -21,8 +23,6 @@ def meta_info() -> Tuple[List[str], List[str]]: return [ "sample_sentence_similarity", - "sample_token_similarity", - "token_similarity", ], ["input_texts", "sample_tokens", "sample_texts", "greedy_tokens"] def __init__( @@ -42,22 +42,23 @@ def __call__( self, dependencies: Dict[str, np.array], texts: List[str], - model: WhiteboxModel, + model: Model, max_new_tokens: int = 100, ) -> Dict[str, np.ndarray]: - device = model.device() - tokenizer = model.tokenizer + if isinstance(model, BlackboxModel): + if torch.cuda.is_available(): + device = "cuda" + else: + device = "cpu" + else: + device = model.device() if not self.crossencoder_setup: self._setup(device=device) self.crossencoder_setup = True - batch_sample_tokens = dependencies["sample_tokens"] batch_texts = dependencies["sample_texts"] batch_input_texts = dependencies["input_texts"] - batch_greedy_tokens = dependencies["greedy_tokens"] - - special_tokens = list(model.tokenizer.added_tokens_decoder.keys()) batch_pairs = [] batch_invs = [] @@ -70,6 +71,71 @@ def __call__( batch_invs.append(inv) batch_counts.append(len(unique_texts)) + sim_matrices = [] + for i, pairs in enumerate(batch_pairs): + sim_scores = self.crossencoder.predict(pairs, batch_size=self.batch_size) + unique_mat_shape = (batch_counts[i], batch_counts[i]) + + sim_scores_matrix = sim_scores.reshape(unique_mat_shape) + inv = batch_invs[i] + + # Recover full matrices from unques by gathering along both axes + # using inverse index + sim_matrices.append(sim_scores_matrix[inv, :][:, inv]) + sim_matrices = np.stack(sim_matrices) + + return { + "sample_sentence_similarity": sim_matrices, + } + +class TokenCrossEncoderSimilarityMatrixCalculator(StatCalculator): + """ + Calculates the cross-encoder similarity matrix for generation samples using RoBERTa model. + """ + + @staticmethod + def meta_info() -> Tuple[List[str], List[str]]: + """ + Returns the statistics and dependencies for the calculator. + """ + + return [ + "sample_token_similarity", + "token_similarity", + ], ["input_texts", "sample_tokens", "greedy_tokens"] + + def __init__( + self, + batch_size: int = 10, + cross_encoder_name: str = "cross-encoder/stsb-roberta-large", + ): + super().__init__() + self.crossencoder_setup = False + self.batch_size = batch_size + self.cross_encoder_name = cross_encoder_name + + def _setup(self, device="cuda"): + self.crossencoder = CrossEncoder(self.cross_encoder_name, device=device) + + def __call__( + self, + dependencies: Dict[str, np.array], + texts: List[str], + model: WhiteboxModel, + max_new_tokens: int = 100, + ) -> Dict[str, np.ndarray]: + device = model.device() + tokenizer = model.tokenizer + special_tokens = list(model.tokenizer.added_tokens_decoder.keys()) + + if not self.crossencoder_setup: + self._setup(device=device) + self.crossencoder_setup = True + + batch_sample_tokens = dependencies["sample_tokens"] + batch_input_texts = dependencies["input_texts"] + batch_greedy_tokens = dependencies["greedy_tokens"] + batch_token_scores = [] for input_texts, tokens in zip(batch_input_texts, batch_greedy_tokens): if len(tokens) > 1: @@ -99,19 +165,6 @@ def __call__( token_scores = np.array([0.5] * len(tokens)) batch_token_scores.append(token_scores) - sim_matrices = [] - for i, pairs in enumerate(batch_pairs): - sim_scores = self.crossencoder.predict(pairs, batch_size=self.batch_size) - unique_mat_shape = (batch_counts[i], batch_counts[i]) - - sim_scores_matrix = sim_scores.reshape(unique_mat_shape) - inv = batch_invs[i] - - # Recover full matrices from unques by gathering along both axes - # using inverse index - sim_matrices.append(sim_scores_matrix[inv, :][:, inv]) - sim_matrices = np.stack(sim_matrices) - batch_samples_token_scores = [] for sample_tokens, input_texts in zip(batch_sample_tokens, batch_input_texts): samples_token_scores = [] @@ -145,7 +198,6 @@ def __call__( batch_samples_token_scores.append(samples_token_scores) return { - "sample_sentence_similarity": sim_matrices, "sample_token_similarity": batch_samples_token_scores, "token_similarity": batch_token_scores, } diff --git a/src/lm_polygraph/stat_calculators/embeddings.py b/src/lm_polygraph/stat_calculators/embeddings.py index 026eb1d58..0f140b2f9 100644 --- a/src/lm_polygraph/stat_calculators/embeddings.py +++ b/src/lm_polygraph/stat_calculators/embeddings.py @@ -4,7 +4,7 @@ from typing import Dict, List, Tuple from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel def get_embeddings_from_output( diff --git a/src/lm_polygraph/stat_calculators/entropy.py b/src/lm_polygraph/stat_calculators/entropy.py index ee04624e1..ad3b21f73 100644 --- a/src/lm_polygraph/stat_calculators/entropy.py +++ b/src/lm_polygraph/stat_calculators/entropy.py @@ -3,7 +3,7 @@ from typing import Dict, List, Tuple from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel class EntropyCalculator(StatCalculator): diff --git a/src/lm_polygraph/stat_calculators/extract_claims.py b/src/lm_polygraph/stat_calculators/extract_claims.py index b1fbc7eba..ab8dac726 100644 --- a/src/lm_polygraph/stat_calculators/extract_claims.py +++ b/src/lm_polygraph/stat_calculators/extract_claims.py @@ -6,7 +6,7 @@ from .stat_calculator import StatCalculator from lm_polygraph.utils.openai_chat import OpenAIChat -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel from .claim_level_prompts import CLAIM_EXTRACTION_PROMPTS, MATCHING_PROMPTS from tqdm import tqdm diff --git a/src/lm_polygraph/stat_calculators/greedy_alternatives_nli.py b/src/lm_polygraph/stat_calculators/greedy_alternatives_nli.py index bb0cfe3b2..583fb39a9 100644 --- a/src/lm_polygraph/stat_calculators/greedy_alternatives_nli.py +++ b/src/lm_polygraph/stat_calculators/greedy_alternatives_nli.py @@ -4,7 +4,7 @@ from typing import Dict, List, Tuple from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel from lm_polygraph.utils.deberta import Deberta from collections import defaultdict import torch.nn as nn diff --git a/src/lm_polygraph/stat_calculators/greedy_cross_encoder_similarity.py b/src/lm_polygraph/stat_calculators/greedy_cross_encoder_similarity.py index 7ce6d8a40..fd5aa9a47 100644 --- a/src/lm_polygraph/stat_calculators/greedy_cross_encoder_similarity.py +++ b/src/lm_polygraph/stat_calculators/greedy_cross_encoder_similarity.py @@ -1,4 +1,5 @@ import numpy as np +import torch import itertools from typing import Dict, List, Tuple @@ -6,7 +7,8 @@ from .stat_calculator import StatCalculator from sentence_transformers import CrossEncoder -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters import * +from lm_polygraph.utils.model import Model class GreedyCrossEncoderSimilarityMatrixCalculator(StatCalculator): @@ -46,10 +48,16 @@ def __call__( self, dependencies: Dict[str, np.array], texts: List[str], - model: WhiteboxModel, + model: Model, max_new_tokens: int = 100, ) -> Dict[str, np.ndarray]: - device = model.device() + if isinstance(model, BlackboxModel): + if torch.cuda.is_available(): + device = "cuda" + else: + device = "cpu" + else: + device = model.device() if not self.crossencoder_setup: self._setup(device=device) diff --git a/src/lm_polygraph/stat_calculators/greedy_lm_probs.py b/src/lm_polygraph/stat_calculators/greedy_lm_probs.py index 0b899f741..0c38847bd 100644 --- a/src/lm_polygraph/stat_calculators/greedy_lm_probs.py +++ b/src/lm_polygraph/stat_calculators/greedy_lm_probs.py @@ -4,7 +4,7 @@ from typing import Dict, List, Tuple from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel class GreedyLMProbsCalculator(StatCalculator): diff --git a/src/lm_polygraph/stat_calculators/greedy_probs_blackbox.py b/src/lm_polygraph/stat_calculators/greedy_probs_blackbox.py index 96f73d679..8c3974888 100644 --- a/src/lm_polygraph/stat_calculators/greedy_probs_blackbox.py +++ b/src/lm_polygraph/stat_calculators/greedy_probs_blackbox.py @@ -3,7 +3,7 @@ from typing import Dict, List, Tuple from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import BlackboxModel +from lm_polygraph.model_adapters.blackbox_model import BlackboxModel class BlackboxGreedyTextsCalculator(StatCalculator): @@ -58,91 +58,36 @@ def __call__( - 'greedy_tokens' (List[List[str]]): tokens of the generated text (greybox only), - 'greedy_tokens_alternatives' (List[List[List[Tuple[str, float]]]]): alternative tokens with logprobs. """ + generation_inputs = dependencies.get("generation_inputs") + inference_inputs = generation_inputs if generation_inputs is not None else texts + if model.supports_logprobs: # Greybox path: generate with logprobs - sequences = model.generate_texts( - input_texts=texts, + output = model.generate( + input_texts=inference_inputs, max_new_tokens=max_new_tokens, - n=1, output_scores=True, top_logprobs=self.top_logprobs, ) - else: - # Blackbox path: generate just the text - sequences = model.generate_texts( - input_texts=texts, - max_new_tokens=max_new_tokens, - n=1, - ) - - # Process the results - greedy_texts = sequences - greedy_log_probs = [] - greedy_log_likelihoods = [] - greedy_tokens = [] - greedy_tokens_alternatives = [] - - # Extract logprobs and tokens from the model's stored data if available - if model.supports_logprobs and hasattr(model, "logprobs") and model.logprobs: - for i, logprob_data in enumerate(model.logprobs): - if hasattr(logprob_data, "content"): - # Extract tokens - tokens = [item.token for item in logprob_data.content] - greedy_tokens.append(tokens) - - # Extract log probabilities for this generation - log_probs_list = [] - log_likelihoods = [] - - # Initialize alternatives structure for this generation - generation_alternatives = [] - - for token_logprobs in logprob_data.content: - # Get the top logprobs for this token position from OpenAI API - token_logprob_dict = {} - - # Build alternatives for this token position - position_alternatives = [] - # First add the chosen token with its logprob - chosen_token = token_logprobs.token - chosen_logprob = token_logprobs.logprob - position_alternatives.append((chosen_token, chosen_logprob)) - - # Then add the remaining top alternatives - for top_logprob in getattr(token_logprobs, "top_logprobs", []): - token_logprob_dict[top_logprob.token] = top_logprob.logprob - # Only add alternatives that are different from the chosen token - if top_logprob.token != chosen_token: - position_alternatives.append( - (top_logprob.token, top_logprob.logprob) - ) - - # Add alternatives for this position to the generation alternatives - generation_alternatives.append(position_alternatives) - - # Create a sparse representation of the logprobs distribution - sparse_logprobs = np.ones(50000) * -float("inf") - - # Map token strings to positions in the sparse array - for token_str, logprob in token_logprob_dict.items(): - token_idx = hash(token_str) % len(sparse_logprobs) - sparse_logprobs[token_idx] = logprob - - log_probs_list.append(sparse_logprobs) - log_likelihoods.append(chosen_logprob) - - greedy_log_probs.append(log_probs_list) - greedy_log_likelihoods.append(log_likelihoods) - greedy_tokens_alternatives.append(generation_alternatives) - - # Ensure all outputs have the same length for greybox case - while len(greedy_tokens) < len(greedy_texts): - # If we're missing token data, add placeholder empty lists - greedy_tokens.append([]) - greedy_log_probs.append([]) - greedy_log_likelihoods.append([]) - greedy_tokens_alternatives.append([]) + # Process the results + greedy_texts = [out.text for out in output] + greedy_tokens = [out.tokens for out in output] + greedy_log_likelihoods = [out.logprobs for out in output] + greedy_log_probs = [np.array(out.top_logprobs) for out in output] + greedy_tokens_alternatives = [] + + # iterate over batch + for out in output: + sample_alternatives = [] + # iterate over generation steps + for i, step_alternatives in enumerate(out.alternative_tokens): + step_alternatives.sort( + key=lambda x: x == out.tokens[i], + reverse=True, + ) + sample_alternatives.append(step_alternatives) + greedy_tokens_alternatives.append(sample_alternatives) result = { "greedy_texts": greedy_texts, @@ -152,8 +97,14 @@ def __call__( "greedy_tokens_alternatives": greedy_tokens_alternatives, } else: + # Blackbox path: generate just the text + output = model.generate( + input_texts=inference_inputs, + max_new_tokens=max_new_tokens, + ) + # For blackbox models, only return the generated texts - result = {"greedy_texts": greedy_texts} + result = {"greedy_texts": [out.text for out in output]} # If user explicitly requests logprobs functionality but model doesn't support it, raise error if any( diff --git a/src/lm_polygraph/stat_calculators/greedy_semantic_matrix.py b/src/lm_polygraph/stat_calculators/greedy_semantic_matrix.py index 07d458d5a..8505714fb 100644 --- a/src/lm_polygraph/stat_calculators/greedy_semantic_matrix.py +++ b/src/lm_polygraph/stat_calculators/greedy_semantic_matrix.py @@ -4,7 +4,7 @@ from typing import Dict, List from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel import torch.nn as nn import torch from tqdm import tqdm diff --git a/src/lm_polygraph/stat_calculators/model_score.py b/src/lm_polygraph/stat_calculators/model_score.py index f5078ae75..25bc15b0c 100644 --- a/src/lm_polygraph/stat_calculators/model_score.py +++ b/src/lm_polygraph/stat_calculators/model_score.py @@ -5,7 +5,7 @@ from torch.nn.utils.rnn import pad_sequence from typing import List, Dict, Tuple from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel log = logging.getLogger(__name__) diff --git a/src/lm_polygraph/stat_calculators/prompt.py b/src/lm_polygraph/stat_calculators/prompt.py index 642ea785f..912b304ff 100644 --- a/src/lm_polygraph/stat_calculators/prompt.py +++ b/src/lm_polygraph/stat_calculators/prompt.py @@ -4,7 +4,7 @@ from typing import Dict, List, Optional, Tuple from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel class BasePromptCalculator(StatCalculator): diff --git a/src/lm_polygraph/stat_calculators/sample_blackbox.py b/src/lm_polygraph/stat_calculators/sample_blackbox.py new file mode 100644 index 000000000..e146dc9d2 --- /dev/null +++ b/src/lm_polygraph/stat_calculators/sample_blackbox.py @@ -0,0 +1,117 @@ +from typing import Dict, List, Tuple + +import numpy as np + +from .stat_calculator import StatCalculator +from lm_polygraph.model_adapters import BlackboxModel + + +class BlackboxSamplingGenerationCalculator(StatCalculator): + """ + Calculates several sampled texts for Blackbox model (lm_polygraph.BlackboxModel). + """ + + @staticmethod + def meta_info() -> Tuple[List[str], List[str]]: + """ + Returns the statistics and dependencies for the calculator. + """ + + return [ + "sample_log_probs", + "sample_tokens", + "sample_texts", + "sample_log_likelihoods", + ], [] + + def __init__(self, samples_n: int = 10): + super().__init__() + self.samples_n = samples_n + + def __call__( + self, + dependencies: Dict[str, np.array], + texts: List[str], + model: BlackboxModel, + max_new_tokens: int = 100, + ) -> Dict[str, np.ndarray]: + """ + Calculates sampled texts for Blackbox model on the input batch. + + Parameters: + dependencies (Dict[str, np.ndarray]): input statistics, can be empty (not used). + texts (List[str]): Input texts batch used for model generation. + model (Model): Model used for generation. + max_new_tokens (int): Maximum number of new tokens at model generation. Default: 100. + Returns: + Dict[str, np.ndarray]: dictionary with the following items: + - 'sample_texts' (List[List[str]]): `samples_n` texts for each input text in the batch, + - 'sample_tokens' (List[List[List[float]]]): tokenized 'sample_texts', + - 'sample_log_probs' (List[List[float]]): sum of the log probabilities at each token of the sampling generation. + - 'sample_log_likelihoods' (List[List[List[float]]]): log probabilities at each token of the sampling generation. + """ + generation_inputs = dependencies.get("generation_inputs") + inference_inputs = generation_inputs if generation_inputs is not None else texts + + if model.supports_logprobs: + # Greybox path: generate with logprobs + output = model.generate_texts( + input_texts=inference_inputs, + max_new_tokens=max_new_tokens, + n=self.samples_n, + output_scores=True, + ) + + sample_texts = [] + sample_log_probs = [] + sample_log_likelihoods = [] + sample_tokens = [] + + # Iterating over batch + for out in output: + sample_texts.append([sample_out.text for sample_out in out]) + sample_tokens.append([sample_out.tokens for sample_out in out]) + sample_log_probs.append( + [np.sum(sample_out.logprobs) for sample_out in out] + ) + sample_log_likelihoods.append( + [sample_out.logprobs for sample_out in out] + ) + + result = { + "sample_texts": sample_texts, + "sample_log_probs": sample_log_probs, + "sample_log_likelihoods": sample_log_likelihoods, + "sample_tokens": sample_tokens, + } + else: + # Blackbox path: generate just the text + output = model.generate_texts( + input_texts=inference_inputs, + max_new_tokens=max_new_tokens, + n=self.samples_n, + ) + + sample_texts = [] + for out in output: + sample_texts.append([sample_out.text for sample_out in out]) + + result = { + "sample_texts": sample_texts, + } + + # If user explicitly requests logprobs functionality but model doesn't support it, raise error + if any( + key in dependencies + for key in [ + "sample_log_probs", + "sample_log_likelihoods", + "sample_tokens", + ] + ): + raise ValueError( + "Model must support logprobs for retrieving probability information. " + "The current model only supports text generation." + ) + + return result diff --git a/src/lm_polygraph/stat_calculators/semantic_classes.py b/src/lm_polygraph/stat_calculators/semantic_classes.py index 4254d4b46..b7446fe5a 100644 --- a/src/lm_polygraph/stat_calculators/semantic_classes.py +++ b/src/lm_polygraph/stat_calculators/semantic_classes.py @@ -4,7 +4,7 @@ from typing import Dict, List, Tuple from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel class SemanticClassesCalculator(StatCalculator): diff --git a/src/lm_polygraph/stat_calculators/semantic_classes_claim_to_samples.py b/src/lm_polygraph/stat_calculators/semantic_classes_claim_to_samples.py index b09af260c..93d3bdfeb 100644 --- a/src/lm_polygraph/stat_calculators/semantic_classes_claim_to_samples.py +++ b/src/lm_polygraph/stat_calculators/semantic_classes_claim_to_samples.py @@ -3,7 +3,7 @@ from typing import Dict, List, Tuple from lm_polygraph.stat_calculators.stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel from lm_polygraph.utils.deberta import Deberta from lm_polygraph.stat_calculators.greedy_alternatives_nli import eval_nli_model diff --git a/src/lm_polygraph/stat_calculators/statistic_extraction.py b/src/lm_polygraph/stat_calculators/statistic_extraction.py index 882c680a5..46de0b06e 100644 --- a/src/lm_polygraph/stat_calculators/statistic_extraction.py +++ b/src/lm_polygraph/stat_calculators/statistic_extraction.py @@ -6,7 +6,7 @@ from typing import Dict, List, Tuple from .stat_calculator import StatCalculator -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel from .greedy_probs import GreedyProbsCalculator diff --git a/src/lm_polygraph/utils/__init__.py b/src/lm_polygraph/utils/__init__.py index f6530a868..677fe34fe 100644 --- a/src/lm_polygraph/utils/__init__.py +++ b/src/lm_polygraph/utils/__init__.py @@ -1,4 +1,4 @@ -from .model import WhiteboxModel, BlackboxModel -from .manager import UEManager -from .estimate_uncertainty import estimate_uncertainty -from .dataset import Dataset +# from .model import WhiteboxModel, BlackboxModel +# from .manager import UEManager +# from .estimate_uncertainty import estimate_uncertainty +# from .dataset import Dataset diff --git a/src/lm_polygraph/utils/dataset.py b/src/lm_polygraph/utils/dataset.py index 6d99437f0..2c0e49467 100644 --- a/src/lm_polygraph/utils/dataset.py +++ b/src/lm_polygraph/utils/dataset.py @@ -5,7 +5,7 @@ from sklearn.model_selection import train_test_split from datasets import load_dataset, Dataset as hf_dataset -from typing import Iterable, Tuple, List, Union, Optional +from typing import Any, Iterable, Tuple, List, Union, Optional class Dataset: @@ -13,7 +13,13 @@ class Dataset: Seq2seq dataset for calculating quality of uncertainty estimation method. """ - def __init__(self, x: List[str], y: List[str], batch_size: int): + def __init__( + self, + x: List[str], + y: List[str], + batch_size: int, + generation_inputs: Optional[List[Any]] = None, + ): """ Parameters: x (List[str]): a list of input texts. @@ -23,18 +29,32 @@ def __init__(self, x: List[str], y: List[str], batch_size: int): self.x = x self.y = y self.batch_size = batch_size + self.generation_inputs = generation_inputs + if self.generation_inputs is not None and len(self.generation_inputs) != len(self.x): + raise ValueError("generation_inputs length must match x length") - def __iter__(self) -> Iterable[Tuple[List[str], List[str]]]: + def __iter__( + self, + ) -> Iterable[ + Tuple[List[str], List[str]] + | Tuple[List[str], List[str], List[Any]] + ]: """ Returns: Iterable[Tuple[List[str], List[str]]]: iterates over batches in dataset, returns list of input texts and list of corresponding output texts. """ for i in range(0, len(self.x), self.batch_size): - yield ( - self.x[i : i + self.batch_size], - self.y[i : i + self.batch_size], - ) + batch_x = self.x[i : i + self.batch_size] + batch_y = self.y[i : i + self.batch_size] + if self.generation_inputs is None: + yield (batch_x, batch_y) + else: + yield ( + batch_x, + batch_y, + self.generation_inputs[i : i + self.batch_size], + ) def __len__(self) -> int: """ @@ -52,6 +72,8 @@ def select(self, indices: List[int]): """ self.x = [self.x[i] for i in indices] self.y = [self.y[i] for i in indices] + if self.generation_inputs is not None: + self.generation_inputs = [self.generation_inputs[i] for i in indices] return self def train_test_split(self, test_size: int, seed: int, split: str = "train"): @@ -74,13 +96,25 @@ def train_test_split(self, test_size: int, seed: int, split: str = "train"): test_size=test_size, random_state=seed, ) + if self.generation_inputs is not None: + gen_train, gen_test = train_test_split( + np.array(self.generation_inputs, dtype=object), + test_size=test_size, + random_state=seed, + ) + else: + gen_train = gen_test = None if split == "train": self.x = X_train.tolist() self.y = y_train.tolist() + if gen_train is not None: + self.generation_inputs = gen_train.tolist() else: self.x = X_test.tolist() self.y = y_test.tolist() + if gen_test is not None: + self.generation_inputs = gen_test.tolist() return ( X_train.tolist(), diff --git a/src/lm_polygraph/utils/estimate_uncertainty.py b/src/lm_polygraph/utils/estimate_uncertainty.py index 8fac44bdf..f0ccc4c1b 100644 --- a/src/lm_polygraph/utils/estimate_uncertainty.py +++ b/src/lm_polygraph/utils/estimate_uncertainty.py @@ -1,7 +1,8 @@ from typing import List, Union from dataclasses import dataclass -from lm_polygraph.utils.model import Model, WhiteboxModel +from lm_polygraph.utils.model import Model +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel from lm_polygraph.model_adapters.visual_whitebox_model import VisualWhiteboxModel from lm_polygraph.estimators.estimator import Estimator from lm_polygraph.utils.manager import UEManager @@ -66,9 +67,9 @@ def estimate_uncertainty( ```python >>> from lm_polygraph import BlackboxModel >>> from lm_polygraph.estimators import EigValLaplacian - >>> model = BlackboxModel.from_openai( - ... 'YOUR_OPENAI_TOKEN', - ... 'gpt-3.5-turbo' + >>> model = BlackboxModel( + ... model_path='gpt-3.5-turbo', + ... api_provider_name='openai' ... ) >>> estimator = EigValLaplacian() >>> estimate_uncertainty(model, estimator, input_text='When did Albert Einstein die?') @@ -87,7 +88,8 @@ def estimate_uncertainty( model, [estimator], available_stat_calculators=register_default_stat_calculators( - model_type + model_type, + model=model, ), # TODO: builder_env_stat_calc=BuilderEnvironmentStatCalculator(model), generation_metrics=[], diff --git a/src/lm_polygraph/utils/factory_estimator.py b/src/lm_polygraph/utils/factory_estimator.py index f4eb9f352..1314228dc 100644 --- a/src/lm_polygraph/utils/factory_estimator.py +++ b/src/lm_polygraph/utils/factory_estimator.py @@ -47,6 +47,7 @@ def load_simple_estimators(name: str, config): PTrueClaim, ClaimConditionedProbabilityClaim, RandomBaselineClaim, + RandomBaseline, FrequencyScoringClaim, TokenSARClaim, FocusClaim, diff --git a/src/lm_polygraph/utils/manager.py b/src/lm_polygraph/utils/manager.py index e1b30a48a..aef527214 100644 --- a/src/lm_polygraph/utils/manager.py +++ b/src/lm_polygraph/utils/manager.py @@ -10,7 +10,8 @@ from tqdm import tqdm from lm_polygraph.utils.dataset import Dataset -from lm_polygraph.utils.model import BlackboxModel, Model +from lm_polygraph.model_adapters.blackbox_model import BlackboxModel +from lm_polygraph.utils.model import Model from lm_polygraph.utils.processor import Processor from lm_polygraph.generation_metrics.generation_metric import GenerationMetric from lm_polygraph.ue_metrics.ue_metric import ( @@ -328,7 +329,13 @@ def estimate( def _process(self, iterable_data, batch_callback): iterable_data = tqdm(self.data) if self.verbose else self.data - for batch_i, (inp_texts, target_texts) in enumerate(iterable_data): + for batch_i, batch in enumerate(iterable_data): + generation_inputs = None + if len(batch) == 3: + inp_texts, target_texts, generation_inputs = batch + else: + inp_texts, target_texts = batch + batch_stats: Dict[str, np.ndarray] = {} for key, val in [ ("input_texts", inp_texts), @@ -337,8 +344,14 @@ def _process(self, iterable_data, batch_callback): self.stats[key] += val batch_stats[key] = val batch_stats["model"] = self.model + if generation_inputs is not None: + batch_stats["generation_inputs"] = generation_inputs - batch_stats = self.calculate(batch_stats, self.stat_calculators, inp_texts) + batch_stats = self.calculate( + batch_stats, + self.stat_calculators, + inp_texts, + ) batch_estimations, bad_estimators = self.estimate( batch_stats, self.estimators diff --git a/src/lm_polygraph/utils/model.py b/src/lm_polygraph/utils/model.py index 735ce9184..d9e981082 100644 --- a/src/lm_polygraph/utils/model.py +++ b/src/lm_polygraph/utils/model.py @@ -1,28 +1,7 @@ -import torch -import openai -import time import logging -import json -from dataclasses import asdict -from typing import List, Dict, Optional, Union +from typing import List from abc import abstractmethod, ABC -from transformers import ( - AutoTokenizer, - AutoModelForSeq2SeqLM, - AutoModelForCausalLM, - AutoConfig, - LogitsProcessorList, - BartForConditionalGeneration, -) -from huggingface_hub import InferenceClient - -from lm_polygraph.utils.generation_parameters import ( - GenerationParameters, - GenerationParametersFactory, -) -from lm_polygraph.utils.ensemble_utils.ensemble_generator import EnsembleGenerationMixin -from lm_polygraph.utils.ensemble_utils.dropout import replace_dropout log = logging.getLogger("lm_polygraph") @@ -69,622 +48,3 @@ def __call__(self, **args): Not implemented for black-box models. """ raise Exception("Not implemented") - - -class BlackboxModel(Model): - """ - Black-box model class. Have no access to model scores and logits. - Currently implemented blackbox models: OpenAI models, Huggingface models. - - Examples: - - ```python - >>> from lm_polygraph import BlackboxModel - >>> model = BlackboxModel.from_openai( - ... 'YOUR_OPENAI_TOKEN', - ... 'gpt-3.5-turbo' - ... ) - ``` - - ```python - >>> from lm_polygraph import BlackboxModel - >>> model = BlackboxModel.from_huggingface( - ... hf_api_token='YOUR_API_TOKEN', - ... hf_model_id='google/t5-large-ssm-nqo' - ... ) - ``` - """ - - def __init__( - self, - openai_api_key: str = None, - model_path: str = None, - hf_api_token: str = None, - generation_parameters: GenerationParameters = GenerationParameters(), - supports_logprobs: bool = False, - ): - """ - Parameters: - openai_api_key (Optional[str]): OpenAI API key if the blackbox model comes from OpenAI. Default: None. - model_path (Optional[str]): Unique model path. Openai model name, if `openai_api_key` is specified, - huggingface path, if `hf_api_token` is specified. Default: None. - hf_api_token (Optional[str]): Huggingface API token if the blackbox model comes from HF. Default: None. - generation_parameters (GenerationParameters): parameters to use in model generation. Default: default parameters. - supports_logprobs (bool): Whether the model supports returning log probabilities. Default: False. - """ - super().__init__(model_path, "Blackbox") - self.generation_parameters = generation_parameters - self.openai_api_key = openai_api_key - self.supports_logprobs = supports_logprobs - - if openai_api_key is not None: - self.openai_api = openai.OpenAI(api_key=openai_api_key) - - self.hf_api_token = hf_api_token - - def _validate_args(self, args): - """ - Validates and adapts arguments for BlackboxModel generation. - - Parameters: - args (dict): The arguments to validate. - - Returns: - dict: Validated and adapted arguments. - """ - args_copy = args.copy() - - # BlackboxModel specific validation - for delete_key in [ - "do_sample", - "min_length", - "top_k", - "repetition_penalty", - "min_new_tokens", - "num_beams", - "allow_newlines", - "stop_strings", - ]: - args_copy.pop(delete_key, None) - - # Map HF argument names to OpenAI/HF API argument names - key_mapping = { - "num_return_sequences": "n", - "max_length": "max_tokens", - "max_new_tokens": "max_tokens", - } - for key, replace_key in key_mapping.items(): - if key in args_copy: - args_copy[replace_key] = args_copy[key] - args_copy.pop(key) - - return args_copy - - def _query(self, payload): - client = InferenceClient(model=self.model_path, token=self.hf_api_token) - response = client.chat_completion(payload) - raw_json = json.dumps(response, indent=2) - return raw_json - - @staticmethod - def from_huggingface(hf_api_token: str, hf_model_id: str, **kwargs): - """ - Initializes a blackbox model from huggingface. - - Parameters: - hf_api_token (Optional[str]): Huggingface API token if the blackbox model comes from HF. Default: None. - hf_model_id (Optional[str]): model path in huggingface. - """ - generation_parameters = kwargs.pop( - "generation_parameters", GenerationParameters() - ) - return BlackboxModel( - hf_api_token=hf_api_token, - model_path=hf_model_id, - generation_parameters=generation_parameters, - ) - - @staticmethod - def from_openai( - openai_api_key: str, model_path: str, supports_logprobs: bool = False, **kwargs - ): - """ - Initializes a blackbox model from OpenAI API. - - Parameters: - openai_api_key (Optional[str]): OpenAI API key. Default: None. - model_path (Optional[str]): model name in OpenAI. - supports_logprobs (bool): Whether the model supports returning log probabilities. Default: False. - """ - generation_parameters = kwargs.pop( - "generation_parameters", GenerationParameters() - ) - return BlackboxModel( - openai_api_key=openai_api_key, - model_path=model_path, - supports_logprobs=supports_logprobs, - generation_parameters=generation_parameters, - ) - - def generate_texts(self, input_texts: List[str], **args) -> List[str]: - """ - Generates a list of model answers using input texts batch. - - Parameters: - input_texts (List[str]): input texts batch. - Return: - List[str]: corresponding model generations. Have the same length as `input_texts`. - """ - # Apply default parameters first, then override with provided args - default_params = asdict(self.generation_parameters) - default_params.update(args) - args = self._validate_args(default_params) - - # Check if we're trying to access features that require logprobs support - if ( - any( - args.get(arg, False) - for arg in [ - "output_scores", - "output_attentions", - "output_hidden_states", - ] - ) - and not self.supports_logprobs - ): - raise Exception("Cannot access logits for blackbox model") - - texts = [] - - if self.openai_api_key is not None: - # Save log probabilities if requested - self.last_response = None - self.logprobs = [] - self.tokens = [] - - # Check if we need to return logprobs - return_logprobs = args.pop("output_scores", False) - logprobs_args = {} - - if return_logprobs and self.supports_logprobs: - logprobs_args["logprobs"] = True - # OpenAI supports returning top logprobs, default to 5 - logprobs_args["top_logprobs"] = args.pop("top_logprobs", 5) - - for prompt in input_texts: - if isinstance(prompt, str): - # If prompt is a string, create a single message with "user" role - messages = [{"role": "user", "content": prompt}] - elif isinstance(prompt, list) and all( - isinstance(item, dict) for item in prompt - ): - # If prompt is a list of dictionaries, assume it's already structured as chat - messages = prompt - else: - raise ValueError( - "Invalid prompt format. Must be either a string or a list of dictionaries." - ) - - retries = 0 - while True: - try: - response = self.openai_api.chat.completions.create( - model=self.model_path, - messages=messages, - **args, - **logprobs_args, - ) - break - except Exception as e: - if retries > 4: - raise Exception from e - else: - retries += 1 - continue - - if args.get("n", 1) == 1: - texts.append(response.choices[0].message.content) - # Store logprobs if available - if return_logprobs and hasattr(response.choices[0], "logprobs"): - self.logprobs.append(response.choices[0].logprobs) - # Extract token information if available - if hasattr(response.choices[0].logprobs, "content"): - tokens = [ - item.token - for item in response.choices[0].logprobs.content - ] - self.tokens.append(tokens) - else: - texts.append([resp.message.content for resp in response.choices]) - # For multiple returns, we don't collect logprobs for now - - # Store the last response for later use - self.last_response = response - - elif (self.hf_api_token is not None) & (self.model_path is not None): - for prompt in input_texts: - start = time.time() - while True: - current_time = time.time() - messages = [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": prompt}, - ] - output = self._query(messages) - - if isinstance(output, dict): - if (list(output.keys())[0] == "error") & ( - "estimated_time" in output.keys() - ): - estimated_time = float(output["estimated_time"]) - elapsed_time = current_time - start - print( - f"{output['error']}. Estimated time: {round(estimated_time - elapsed_time, 2)} sec." - ) - time.sleep(5) - elif (list(output.keys())[0] == "error") & ( - "estimated_time" not in output.keys() - ): - log.error(f"{output['error']}") - break - elif isinstance(output, list): - break - - texts.append(output[0]["generated_text"]) - else: - print( - "Please provide HF API token and model id for using models from HF or openai API key for using OpenAI models" - ) - - return texts - - def generate(self, **args): - """ - For OpenAI models with logprobs support, returns a lightweight wrapper around OpenAI API response. - For other blackbox models, raises an exception as this is not implemented. - - Parameters: - **args: Arguments to pass to the generate method. - Returns: - object: A wrapper around the OpenAI API response if logprobs are supported. - Raises: - Exception: If the model doesn't support logprobs. - """ - if self.supports_logprobs: - # Apply default parameters first, then override with provided args - default_params = asdict(self.generation_parameters) - default_params.update(args) - args = self._validate_args(default_params) - - args["output_scores"] = True - sequences = self.generate_texts(**args) - - # Return a simple object with the necessary attributes for compatibility - class OpenAIGenerationOutput: - def __init__(self, sequences, scores): - self.sequences = sequences - self.scores = scores - - return OpenAIGenerationOutput(sequences, self.logprobs) - else: - raise Exception("Cannot access logits of blackbox model") - - def __call__(self, **args): - """ - Not implemented for blackbox models. - """ - raise Exception("Cannot access logits of blackbox model") - - def tokenizer(self, *args, **kwargs): - """ - Not implemented for blackbox models. - """ - raise Exception("Cannot access logits of blackbox model") - - -class WhiteboxModel(Model): - """ - White-box model class. Have access to model scores and logits. Currently implemented only for Huggingface models. - - Examples: - - ```python - >>> from lm_polygraph import WhiteboxModel - >>> model = WhiteboxModel.from_pretrained( - ... "bigscience/bloomz-3b", - ... ) - ``` - """ - - def __init__( - self, - model: AutoModelForCausalLM, - tokenizer: AutoTokenizer, - model_path: str = None, - model_type: str = "CausalLM", - generation_parameters: GenerationParameters = GenerationParameters(), - instruct: bool = False, - ): - """ - Parameters: - model (AutoModelForCausalLM): HuggingFace model. - tokenizer (AutoTokenizer): HuggingFace tokenizer. - model_path (Optional[str]): Unique model path in HuggingFace. - model_type (str): Additional model specifications. - parameters (GenerationParameters): parameters to use in model generation. Default: default parameters. - """ - super().__init__(model_path, model_type) - self.model = model - self.tokenizer = tokenizer - self.generation_parameters = generation_parameters - self.instruct = instruct - - def _validate_args(self, args): - """ - Validates and adapts arguments for WhiteboxModel generation. - - Parameters: - args (dict): The arguments to validate. - - Returns: - dict: Validated and adapted arguments. - """ - args_copy = args.copy() - - # WhiteboxModel specific validation - if "presence_penalty" in args_copy and args_copy["presence_penalty"] != 0.0: - log.warning( - "Skipping requested argument presence_penalty={}".format( - args_copy["presence_penalty"] - ) - ) - - # Remove arguments that are not supported by the HF model.generate function - keys_to_remove = ["presence_penalty", "allow_newlines"] - for key in keys_to_remove: - args_copy.pop(key, None) - - return args_copy - - class _ScoresProcessor: - # Stores original token scores instead of the ones modified with generation parameters - def __init__(self): - self.scores = [] - - def __call__(self, input_ids=None, scores=None): - self.scores.append(scores.log_softmax(-1)) - return scores - - def generate(self, **args): - """ - Generates the model output with scores from batch formed by HF Tokenizer. - - Parameters: - **args: Any arguments that can be passed to model.generate function from HuggingFace. - Returns: - ModelOutput: HuggingFace generation output with scores overriden with original probabilities. - """ - default_params = asdict(self.generation_parameters) - - # add ScoresProcessor to collect original scores - processor = self._ScoresProcessor() - if "logits_processor" in args.keys(): - logits_processor = LogitsProcessorList( - [processor, args["logits_processor"]] - ) - else: - logits_processor = LogitsProcessorList([processor]) - args["logits_processor"] = logits_processor - - # update default parameters with passed arguments - default_params.update(args) - args = default_params - - if "stop_strings" in args: - args["tokenizer"] = self.tokenizer - - args = self._validate_args(args) - generation = self.model.generate(**args) - - # override generation.scores with original scores from model - generation.generation_scores = generation.scores - generation.scores = processor.scores - - return generation - - def generate_texts(self, input_texts: List[str], **args) -> List[str]: - """ - Generates a list of model answers using input texts batch. - - Parameters: - input_texts (List[str]): input texts batch. - Return: - List[str]: corresponding model generations. Have the same length as `input_texts`. - """ - # Apply default parameters first, then override with provided args - default_params = asdict(self.generation_parameters) - default_params.update(args) - args = self._validate_args(default_params) - - args["return_dict_in_generate"] = True - batch: Dict[str, torch.Tensor] = self.tokenize(input_texts) - batch = {k: v.to(self.device()) for k, v in batch.items()} - sequences = self.generate(**batch, **args).sequences.cpu() - input_len = batch["input_ids"].shape[1] - texts = [] - - decode_args = {} - if self.tokenizer.chat_template is not None: - decode_args["skip_special_tokens"] = True - - for seq in sequences: - if self.model_type == "CausalLM": - texts.append(self.tokenizer.decode(seq[input_len:], **decode_args)) - else: - texts.append(self.tokenizer.decode(seq[1:], **decode_args)) - - return texts - - def __call__(self, **args): - """ - Calls the model on the input batch. Returns the resulted scores. - """ - return self.model(**args) - - def device(self): - """ - Returns the device the model is currently loaded on. - - Returns: - str: device string. - """ - return self.model.device - - @staticmethod - def from_pretrained( - model_path: str, - generation_params: Optional[Dict] = {}, - add_bos_token: bool = True, - **kwargs, - ): - """ - Initializes the model from HuggingFace. Automatically determines model type. - - Parameters: - model_path (str): model path in HuggingFace. - generation_params (Dict): generation arguments for - lm_polygraph.utils.generation_parametersGenerationParameters - add_bos_token (bool): tokenizer argument. Default: True. - """ - log.warning( - "WhiteboxModel#from_pretrained is deprecated and will be removed in the next release. Please instantiate WhiteboxModel directly by passing an already loaded model, tokenizer and model path." - ) - - config = AutoConfig.from_pretrained( - model_path, trust_remote_code=True, **kwargs - ) - - if any(["CausalLM" in architecture for architecture in config.architectures]): - model_type = "CausalLM" - model = AutoModelForCausalLM.from_pretrained( - model_path, trust_remote_code=True, **kwargs - ) - elif any( - [ - ("Seq2SeqLM" in architecture) - or ("ConditionalGeneration" in architecture) - for architecture in config.architectures - ] - ): - model_type = "Seq2SeqLM" - model = AutoModelForSeq2SeqLM.from_pretrained(model_path, **kwargs) - if "falcon" in model_path: - model.transformer.alibi = True - elif any( - ["JAISLMHeadModel" in architecture for architecture in config.architectures] - ): - model_type = "CausalLM" - model = AutoModelForCausalLM.from_pretrained( - model_path, - trust_remote_code=True, - **kwargs, - ) - elif any( - ["BartModel" in architecture for architecture in config.architectures] - ): - model_type = "Seq2SeqLM" - model = BartForConditionalGeneration.from_pretrained(model_path, **kwargs) - else: - raise ValueError( - f"Model {model_path} is not adapted for the sequence generation task" - ) - - tokenizer = AutoTokenizer.from_pretrained( - model_path, - padding_side="left", - add_bos_token=add_bos_token, - **kwargs, - ) - - model.eval() - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - - generation_params = GenerationParametersFactory.from_params( - yaml_config=generation_params, - native_config=asdict(model.config), - ) - - instance = WhiteboxModel( - model, tokenizer, model_path, model_type, generation_params - ) - - return instance - - def tokenize( - self, texts: Union[List[str], List[List[Dict[str, str]]]] - ) -> Dict[str, torch.Tensor]: - """ - Tokenizes input texts batch into a dictionary using the model tokenizer. - - Parameters: - texts (List[str]): list of input texts batch. - Returns: - dict[str, torch.Tensor]: tensors dictionary obtained by tokenizing input texts batch. - """ - # Apply chat template if tokenizer has it - add_start_symbol = True - if self.instruct: - formatted_texts = [] - for chat in texts: - if isinstance(chat, str): - chat = [{"role": "user", "content": chat}] - formatted_chat = self.tokenizer.apply_chat_template( - chat, add_generation_prompt=True, tokenize=False - ) - formatted_texts.append(formatted_chat) - texts = formatted_texts - - add_start_symbol = False - return self.tokenizer( - texts, - padding=True, - return_tensors="pt", - add_special_tokens=add_start_symbol, - ) - - -def create_ensemble( - models: List[WhiteboxModel] = [], - mc: bool = False, - seed: int = 1, - mc_seeds: List[int] = [1], - ensembling_mode: str = "pe", - dropout_rate: float = 0.1, - **kwargs, -) -> WhiteboxModel: - model = models[0] - ens = model.model - - ens.__class__ = type( - "EnsembleModel", (model.model.__class__, EnsembleGenerationMixin), {} - ) - - if mc: - ens.mc = True - ens.mc_seeds = mc_seeds - ens.base_seed = seed - ens.ensembling_mode = ensembling_mode - ens.mc_models_num = len(mc_seeds) - ens.mc_seeds = mc_seeds - - replace_dropout( - ens.config._name_or_path, ens, p=dropout_rate, share_across_tokens=True - ) - ens.train() - else: - raise ValueError( - "Only Monte-Carlo ensembling is available. Please set the corresponding argument value to True" - ) - - return model diff --git a/test/configs/model/gpt-4o-mini-greybox.yaml b/test/configs/model/gpt-4o-mini-greybox.yaml index b3898f364..f859dbfa4 100644 --- a/test/configs/model/gpt-4o-mini-greybox.yaml +++ b/test/configs/model/gpt-4o-mini-greybox.yaml @@ -4,4 +4,3 @@ defaults: provider: openai path: gpt-4o-mini type: Blackbox -supports_logprobs: true diff --git a/test/local/test_openai_adapter.py b/test/local/test_openai_adapter.py new file mode 100644 index 000000000..b6926662d --- /dev/null +++ b/test/local/test_openai_adapter.py @@ -0,0 +1,110 @@ +""" +Smoke test suite for OpenAI API integration with lm-polygraph + +This test mirrors the Together.ai adapter tests to ensure the OpenAI adapter +works correctly with registration, request adaptation, response parsing, and +live API calls when an `OPENAI_API_KEY` is available. + +Usage: + # Set your OpenAI API key + export OPENAI_API_KEY="your-api-key-here" + + # Run the tests + pytest test/local/test_openai_adapter.py -v + + # Skip smoke tests if no API key is set + pytest test/local/test_openai_adapter.py -v -k "not openai_api_key" +""" + +import os + +import pytest + +# Import the local OpenAI adapter so it registers itself +from lm_polygraph.model_adapters.openai_adapter import OpenAIAdapter + +from lm_polygraph.model_adapters.blackbox_model import BlackboxModel +from lm_polygraph.estimators import Perplexity +from lm_polygraph.utils.estimate_uncertainty import estimate_uncertainty +from lm_polygraph.model_adapters.api_provider_adapter import ( + get_adapter, +) + +# Load environment variables from .env file +try: + from dotenv import load_dotenv + + load_dotenv() +except ImportError: + print("Warning: python-dotenv not installed. Using environment variables only.") + + +@pytest.fixture(scope="module") +def openai_api_key(): + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + pytest.skip("OPENAI_API_KEY environment variable not set") + return api_key + + +def test_openai_adapter_registration(): + """Ensure the OpenAI adapter is registered and retrievable.""" + adapter = get_adapter("openai") + assert adapter is not None + assert adapter.__class__.__name__ == "OpenAIAdapter" + + +def test_openai_adapter_functionality(): + """Directly test adapter parameter adaptation and validation.""" + adapter = OpenAIAdapter() + + test_params = { + "max_new_tokens": 50, + "min_length": 10, + "top_k": 40, + "repetition_penalty": 1.2, + "num_beams": 3, + "temperature": 0.7, + "top_p": 0.9, + "do_sample": True, + "top_logprobs": 3, + "stop_strings": ["\n\n"], + "output_scores": True, + } + + adapted = adapter.adapt_request(test_params) + + assert "min_length" not in adapted, "min_length should be removed" + assert "top_k" not in adapted, "top_k should be removed" + assert "repetition_penalty" not in adapted, "repetition_penalty should be removed" + assert "num_beams" not in adapted, "num_beams should be removed" + assert "do_sample" not in adapted, "do_sample should be removed" + assert "stop_strings" not in adapted, "stop_strings should be removed" + assert "max_completion_tokens" in adapted, "max_new_tokens should be mapped to max_tokens" + assert adapted["max_completion_tokens"] == 50 + assert adapted["top_logprobs"] == 3 + assert adapted["output_scores"] is True + assert adapted["temperature"] == 0.7 + + params_with_invalid_ranges = { + "temperature": 3.5, + "top_p": 1.5, + "presence_penalty": 3.0, + "frequency_penalty": -3.0, + } + + validated = adapter.validate_parameter_ranges(params_with_invalid_ranges) + assert validated["temperature"] == 2.0 + assert validated["top_p"] == 1.0 + assert validated["presence_penalty"] == 2.0 + assert validated["frequency_penalty"] == -2.0 + + assert adapter.supports_logprobs(), "OpenAI should report logprobs support" + + +def test_openai_adapter_parse_response_error(): + """Invalid response structures should raise ValueError.""" + adapter = OpenAIAdapter() + + with pytest.raises(ValueError): + adapter.parse_response({"unexpected": "structure"}) diff --git a/test/local/test_togetherai_adapter.py b/test/local/test_togetherai_adapter.py new file mode 100644 index 000000000..b2e178872 --- /dev/null +++ b/test/local/test_togetherai_adapter.py @@ -0,0 +1,246 @@ +""" +Smoke test for Together.ai API integration with lm-polygraph + +This test verifies that the together.ai adapter works correctly with actual API calls +using the estimate_uncertainty method and Perplexity estimator. + +Usage: + # Set your Together.ai API key + export OPENAI_API_KEY="your-api-key-here" + + # Run the test + pytest test/local/test_together_ai_smoke.py -v + + # Skip if no API key is set + pytest test/local/test_together_ai_smoke.py -v -k "not together_api_key" +""" + +import os + +import pytest + +# Import the local together.ai adapter (this registers it) +import lm_polygraph.model_adapters.togetherai_adapter # noqa: F401 + +from lm_polygraph.model_adapters.blackbox_model import BlackboxModel +from lm_polygraph.estimators import Perplexity +from lm_polygraph.utils.estimate_uncertainty import estimate_uncertainty +from lm_polygraph.model_adapters.api_provider_adapter import ( + get_adapter, + list_available_adapters, +) + +# Load environment variables from .env file +try: + from dotenv import load_dotenv + + load_dotenv() +except ImportError: + print("Warning: python-dotenv not installed. Using environment variables only.") + + +def test_together_ai_adapter_registration(): + """ + Test that the together.ai adapter is properly registered. + """ + available_adapters = list_available_adapters() + assert ( + "together_ai" in available_adapters + ), f"together_ai adapter not registered. Available: {available_adapters}" + + # Test that we can get the adapter + adapter = get_adapter("together_ai") + assert adapter is not None + assert adapter.__class__.__name__ == "TogetherAIAdapter" + + +def test_together_ai_api_smoke(): + """ + Smoke test for together.ai API using estimate_uncertainty with Perplexity estimator. + + This test verifies: + 1. Together.ai adapter correctly formats requests + 2. API call succeeds and returns valid response + 3. Perplexity estimator can process the response + 4. estimate_uncertainty returns expected structure + """ + # Skip test if no API key is provided + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + pytest.skip("OPENAI_API_KEY environment variable not set") + + # Test configuration + test_input = "What is the capital of France?" + model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo" + + # Initialize model with together.ai configuration + model = BlackboxModel( + model_path=model_name, + api_provider_name="together_ai", + ) + + # Override base URL for together.ai (BlackboxModel uses openai client) + if hasattr(model, "openai_api") and model.openai_api: + model.openai_api.base_url = "https://api.together.xyz/v1" + + # Use minimal generation parameters for faster test + model.generation_parameters.max_new_tokens = 20 + model.generation_parameters.temperature = 0.0 # Deterministic for testing + + # Initialize Perplexity estimator + estimator = Perplexity() + + # Run uncertainty estimation + result = estimate_uncertainty( + model=model, estimator=estimator, input_text=test_input + ) + + # Validate result structure + assert result is not None, "estimate_uncertainty returned None" + assert hasattr(result, "uncertainty"), "Result missing uncertainty field" + assert hasattr(result, "input_text"), "Result missing input_text field" + assert hasattr(result, "generation_text"), "Result missing generation_text field" + assert hasattr(result, "model_path"), "Result missing model_path field" + assert hasattr(result, "estimator"), "Result missing estimator field" + + # Validate result values + assert result.input_text == test_input, f"Input text mismatch: {result.input_text}" + assert result.model_path == model_name, f"Model path mismatch: {result.model_path}" + assert result.estimator == str(estimator), f"Estimator mismatch: {result.estimator}" + + # Validate generation text is not empty + assert result.generation_text is not None, "Generation text is None" + assert len(result.generation_text.strip()) > 0, "Generation text is empty" + + # Validate uncertainty is a valid number + assert isinstance( + result.uncertainty, (int, float) + ), f"Uncertainty is not numeric: {type(result.uncertainty)}" + assert not ( + result.uncertainty != result.uncertainty + ), "Uncertainty is NaN" # NaN check + + # Basic sanity checks for perplexity + # Perplexity should be positive (though can be very small) + assert ( + result.uncertainty >= 0 + ), f"Perplexity should be non-negative, got: {result.uncertainty}" + + # Log results for manual verification + print("\n--- Together.ai Smoke Test Results ---") + print(f"Input: {result.input_text}") + print(f"Generated: {result.generation_text}") + print(f"Uncertainty: {result.uncertainty}") + print(f"Model: {result.model_path}") + print(f"Estimator: {result.estimator}") + print(f"Generation tokens: {result.generation_tokens}") + + +def test_together_ai_api_smoke_with_logprobs(): + """ + Enhanced smoke test that specifically tests logprobs functionality with 5 alternatives. + """ + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + pytest.skip("OPENAI_API_KEY environment variable not set") + + test_input = "The sky is" + model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo" + + # Initialize model with logprobs enabled + model = BlackboxModel( + model_path=model_name, + api_provider_name="together_ai", + ) + + # Override base URL for together.ai + if hasattr(model, "openai_api") and model.openai_api: + model.openai_api.base_url = "https://api.together.xyz/v1" + + # Configure for logprobs + model.generation_parameters.max_new_tokens = 5 + model.generation_parameters.temperature = 0.0 + + # Initialize Perplexity estimator (requires logprobs) + estimator = Perplexity() + + # Run uncertainty estimation + result = estimate_uncertainty( + model=model, estimator=estimator, input_text=test_input + ) + + # Validate that we got a valid result + assert result is not None + assert result.generation_text is not None + assert len(result.generation_text.strip()) > 0 + + # Perplexity requires logprobs, so this should work + assert isinstance(result.uncertainty, (int, float)) + assert result.uncertainty >= 0 + + print("\n--- Together.ai Logprobs Test Results ---") + print(f"Input: {result.input_text}") + print(f"Generated: {result.generation_text}") + print(f"Perplexity: {result.uncertainty}") + + +def test_together_ai_adapter_functionality(): + """ + Test the together.ai adapter functionality directly. + """ + # Import and test the adapter directly + from lm_polygraph.model_adapters.togetherai_adapter import TogetherAIAdapter + + adapter = TogetherAIAdapter() + + # Test request adaptation + test_params = { + "max_new_tokens": 50, + "temperature": 0.7, + "top_p": 0.9, + "do_sample": True, # Should be removed + "logprobs": 5, + } + + adapted = adapter.adapt_request(test_params) + + # Check that parameters were adapted correctly + assert "max_tokens" in adapted, "max_new_tokens should be mapped to max_tokens" + assert adapted["max_tokens"] == 50 + assert "do_sample" not in adapted, "do_sample should be removed" + assert adapted["logprobs"] == 5, "logprobs should be preserved" + + # Test parameter validation + test_params_with_invalid = { + "temperature": 3.0, # Should be clamped to 2.0 + "top_k": 500, # Should be clamped to 200 + } + + validated = adapter.validate_parameter_ranges(test_params_with_invalid) + assert validated["temperature"] == 2.0, "Temperature should be clamped to 2.0" + assert validated["top_k"] == 200, "top_k should be clamped to 200" + + # Test logprobs support + assert adapter.supports_logprobs(), "together.ai should support logprobs" + + +def test_together_ai_api_error_handling(): + """ + Test error handling for invalid configurations. + """ + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + pytest.skip("OPENAI_API_KEY environment variable not set") + + # Test with invalid model name + with pytest.raises(Exception): + model = BlackboxModel( + model_path="nonexistent-model-12345", + api_provider_name="together_ai", + ) + + if hasattr(model, "openai_api") and model.openai_api: + model.openai_api.base_url = "https://api.together.xyz/v1" + + estimator = Perplexity() + estimate_uncertainty(model, estimator, "Test input") diff --git a/test/test_blackbox_estimators.py b/test/test_blackbox_estimators.py new file mode 100644 index 000000000..3f6d7e40f --- /dev/null +++ b/test/test_blackbox_estimators.py @@ -0,0 +1,174 @@ +import os +import torch +import pytest + +from lm_polygraph import estimate_uncertainty +from lm_polygraph.estimators import * +from lm_polygraph.model_adapters.blackbox_model import BlackboxModel + +INPUT = "When was Julius Caesar born?" + + +@pytest.fixture(scope="module") +def model(): + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + pytest.skip("OPENAI_API_KEY environment variable not set") + + return BlackboxModel('gpt-4.1-nano') + + +def test_maximum_sequence_probability(model): + estimator = MaximumSequenceProbability() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_perplexity(model): + estimator = Perplexity() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_mean_token_entropy(model): + estimator = MeanTokenEntropy() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_monte_carlo_sequence_entropy(model): + estimator = MonteCarloSequenceEntropy() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_monte_carlo_normalized_sequence_entropy(model): + estimator = MonteCarloNormalizedSequenceEntropy() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_lexical_similarity_rouge1(model): + estimator = LexicalSimilarity(metric="rouge1") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_lexical_similarity_rouge2(model): + estimator = LexicalSimilarity(metric="rouge2") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_lexical_similarity_rougel(model): + estimator = LexicalSimilarity(metric="rougeL") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_lexical_similarity_bleu(model): + estimator = LexicalSimilarity(metric="BLEU") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_num_sem_sets(model): + estimator = NumSemSets() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_eigval_laplacian_nli_entail(model): + estimator = EigValLaplacian(similarity_score="NLI_score", affinity="entail") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_eigval_laplacian_nli_contra(model): + estimator = EigValLaplacian(similarity_score="NLI_score", affinity="contra") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_eigval_laplacian_jaccard(model): + estimator = EigValLaplacian(similarity_score="Jaccard_score") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_degmat_nli_entail(model): + estimator = DegMat(similarity_score="NLI_score", affinity="entail") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_degmat_nli_contra(model): + estimator = DegMat(similarity_score="NLI_score", affinity="contra") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_degmat_jaccard(model): + estimator = DegMat(similarity_score="Jaccard_score") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_eccentricity_nli_entail(model): + estimator = Eccentricity(similarity_score="NLI_score", affinity="entail") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_eccentricity_nli_contra(model): + estimator = Eccentricity(similarity_score="NLI_score", affinity="contra") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_eccentricity_jaccard(model): + estimator = Eccentricity(similarity_score="Jaccard_score") + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_semantic_entropy(model): + estimator = SemanticEntropy() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_sentence_sar(model): + estimator = SentenceSAR() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_cocoamsp(model): + estimator = CocoaMSP() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_cocoappl(model): + estimator = CocoaPPL() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_cocoamte(model): + estimator = CocoaMTE() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_semantic_density_concat(model): + estimator = SemanticDensity() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + +def test_semantic_density(model): + estimator = SemanticDensity(concat_input=False) + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) diff --git a/test/test_blackbox_generation_inputs_routing.py b/test/test_blackbox_generation_inputs_routing.py new file mode 100644 index 000000000..4b6026003 --- /dev/null +++ b/test/test_blackbox_generation_inputs_routing.py @@ -0,0 +1,57 @@ +from lm_polygraph.stat_calculators.greedy_probs_blackbox import ( + BlackboxGreedyTextsCalculator, +) +from lm_polygraph.stat_calculators.sample_blackbox import ( + BlackboxSamplingGenerationCalculator, +) + + +class _DummyOutput: + def __init__(self, text: str): + self.text = text + + +class _DummyBlackboxModel: + supports_logprobs = False + + def __init__(self): + self.last_input_texts = None + + def generate_texts(self, input_texts, max_new_tokens=100, n=1, output_scores=False): + self.last_input_texts = input_texts + if n == 1: + return [_DummyOutput("greedy") for _ in input_texts] + return [[_DummyOutput(f"sample_{i}") for i in range(n)] for _ in input_texts] + + +def test_blackbox_greedy_calculator_uses_generation_inputs_when_provided(): + model = _DummyBlackboxModel() + calc = BlackboxGreedyTextsCalculator() + generation_inputs = [[{"role": "user", "content": "multimodal"}]] + + result = calc( + {"generation_inputs": generation_inputs}, + texts=["plain-text-prompt"], + model=model, + max_new_tokens=5, + ) + + assert model.last_input_texts == generation_inputs + assert result["greedy_texts"] == ["greedy"] + + +def test_blackbox_sampling_calculator_uses_generation_inputs_when_provided(): + model = _DummyBlackboxModel() + calc = BlackboxSamplingGenerationCalculator(samples_n=3) + generation_inputs = [[{"role": "user", "content": "multimodal"}]] + + result = calc( + {"generation_inputs": generation_inputs}, + texts=["plain-text-prompt"], + model=model, + max_new_tokens=5, + ) + + assert model.last_input_texts == generation_inputs + assert result["sample_texts"] == [["sample_0", "sample_1", "sample_2"]] + diff --git a/test/test_blackbox_model_logprobs_override.py b/test/test_blackbox_model_logprobs_override.py new file mode 100644 index 000000000..bdb63bb08 --- /dev/null +++ b/test/test_blackbox_model_logprobs_override.py @@ -0,0 +1,99 @@ +import pytest + +from lm_polygraph.model_adapters.blackbox_model import BlackboxModel + + +class _DummyResponse: + def __init__(self, text: str): + self.text = text + + +class _DummyAdapter: + def __init__(self, supports_logprobs: bool): + self._supports_logprobs = supports_logprobs + self.supports_logprobs_calls = 0 + self.generate_calls = 0 + self.last_generate_args = None + + def supports_logprobs(self, model_path=None): + self.supports_logprobs_calls += 1 + return self._supports_logprobs + + def validate_parameter_ranges(self, params: dict) -> dict: + return params + + def adapt_request(self, params: dict) -> dict: + return params + + def generate_texts(self, model, input_texts, args): + self.generate_calls += 1 + self.last_generate_args = args + return [[_DummyResponse("ok")] for _ in input_texts] + + +def test_blackbox_model_uses_adapter_support_by_default(monkeypatch): + adapter = _DummyAdapter(supports_logprobs=False) + monkeypatch.setattr( + "lm_polygraph.model_adapters.blackbox_model.get_adapter", lambda _: adapter + ) + + model = BlackboxModel(model_path="dummy-model", api_provider_name="openai") + + assert model.supports_logprobs is False + assert adapter.supports_logprobs_calls == 1 + + +def test_blackbox_model_override_supports_logprobs_true(monkeypatch): + adapter = _DummyAdapter(supports_logprobs=False) + monkeypatch.setattr( + "lm_polygraph.model_adapters.blackbox_model.get_adapter", lambda _: adapter + ) + + model = BlackboxModel( + model_path="dummy-model", + api_provider_name="openai", + supports_logprobs=True, + ) + + assert model.supports_logprobs is True + assert adapter.supports_logprobs_calls == 0 + + +def test_blackbox_model_override_supports_logprobs_false_blocks_output_scores( + monkeypatch, +): + adapter = _DummyAdapter(supports_logprobs=True) + monkeypatch.setattr( + "lm_polygraph.model_adapters.blackbox_model.get_adapter", lambda _: adapter + ) + + model = BlackboxModel( + model_path="dummy-model", + api_provider_name="openai", + supports_logprobs=False, + ) + + with pytest.raises(Exception, match="Cannot access logits"): + model.generate_texts(input_texts=["prompt"], output_scores=True) + + assert adapter.generate_calls == 0 + + +def test_blackbox_model_override_supports_logprobs_true_allows_output_scores( + monkeypatch, +): + adapter = _DummyAdapter(supports_logprobs=False) + monkeypatch.setattr( + "lm_polygraph.model_adapters.blackbox_model.get_adapter", lambda _: adapter + ) + + model = BlackboxModel( + model_path="dummy-model", + api_provider_name="openai", + supports_logprobs=True, + ) + output = model.generate_texts(input_texts=["prompt"], output_scores=True) + + assert adapter.generate_calls == 1 + assert adapter.last_generate_args["output_scores"] is True + assert output[0][0].text == "ok" diff --git a/test/test_dataset_generation_inputs.py b/test/test_dataset_generation_inputs.py new file mode 100644 index 000000000..1eb417538 --- /dev/null +++ b/test/test_dataset_generation_inputs.py @@ -0,0 +1,50 @@ +import pytest + +from lm_polygraph.utils.dataset import Dataset + + +def test_dataset_iter_yields_generation_inputs_when_provided(): + dataset = Dataset( + x=["p1", "p2", "p3"], + y=["", "", ""], + batch_size=2, + generation_inputs=["g1", "g2", "g3"], + ) + + batches = list(dataset) + assert len(batches) == 2 + + x1, y1, g1 = batches[0] + assert x1 == ["p1", "p2"] + assert y1 == ["", ""] + assert g1 == ["g1", "g2"] + + x2, y2, g2 = batches[1] + assert x2 == ["p3"] + assert y2 == [""] + assert g2 == ["g3"] + + +def test_dataset_select_keeps_generation_inputs_aligned(): + dataset = Dataset( + x=["p1", "p2", "p3"], + y=["t1", "t2", "t3"], + batch_size=2, + generation_inputs=["g1", "g2", "g3"], + ) + dataset.select([2, 0]) + + assert dataset.x == ["p3", "p1"] + assert dataset.y == ["t3", "t1"] + assert dataset.generation_inputs == ["g3", "g1"] + + +def test_dataset_generation_inputs_length_validation(): + with pytest.raises(ValueError, match="generation_inputs length must match x length"): + Dataset( + x=["p1", "p2"], + y=["", ""], + batch_size=1, + generation_inputs=["g1"], + ) + diff --git a/test/test_estimators.py b/test/test_estimators.py index 9bc68bdf8..bd5d3223d 100644 --- a/test/test_estimators.py +++ b/test/test_estimators.py @@ -5,7 +5,7 @@ from lm_polygraph import estimate_uncertainty from lm_polygraph.estimators import * -from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel INPUT = "When was Julius Caesar born?" diff --git a/test/test_random_baseline_estimator.py b/test/test_random_baseline_estimator.py new file mode 100644 index 000000000..85ba81829 --- /dev/null +++ b/test/test_random_baseline_estimator.py @@ -0,0 +1,49 @@ +import numpy as np + +from lm_polygraph.defaults.register_default_stat_calculators import ( + register_default_stat_calculators, +) +from lm_polygraph.estimators import RandomBaseline +from lm_polygraph.utils.manager import order_calculators +from lm_polygraph.utils.factory_estimator import FactoryEstimator + + +def test_random_baseline_scores_shape_and_range(): + estimator = RandomBaseline() + scores = estimator({"input_texts": ["a", "b", "c"]}) + + assert scores.shape == (3,) + assert np.all(scores >= 0.0) + assert np.all(scores <= 1.0) + + +def test_random_baseline_has_no_stat_dependencies(): + estimator = RandomBaseline() + assert estimator.stats_dependencies == [] + assert estimator.level == "sequence" + + +def test_random_baseline_available_in_factory(): + estimator = FactoryEstimator()("RandomBaseline", {}) + assert isinstance(estimator, RandomBaseline) + + +def test_random_baseline_resolves_to_greedy_probs_calculator_only(): + estimator = RandomBaseline() + calculators = register_default_stat_calculators(model_type="Whitebox") + + stat_calculators = {} + stat_dependencies = {} + for calculator in calculators: + for stat in calculator.stats: + stat_calculators[stat] = calculator + stat_dependencies[stat] = calculator.dependencies + + ordered_stats, _ = order_calculators( + estimator.stats_dependencies + ["greedy_texts", "greedy_tokens"], + stat_calculators, + stat_dependencies, + ) + + assert ordered_stats == ["greedy_texts"] + assert stat_calculators[ordered_stats[0]].name == "GreedyProbsCalculator"