uqlm.nli.nli.NLI#

class uqlm.nli.nli.NLI(device=None, verbose=False, nli_model_name='microsoft/deberta-large-mnli', max_length=2000, batch_size=32, use_fp16=False, device_map=None)#

Bases: object

__init__(device=None, verbose=False, nli_model_name='microsoft/deberta-large-mnli', max_length=2000, batch_size=32, use_fp16=False, device_map=None)#

A class to computing NLI-based confidence scores. This class offers two types of confidence scores, namely noncontradiction probability [1] and semantic entropy [2].

Parameters:

  • device (torch.device input or torch.device object, default=None) – Specifies the device that classifiers use for prediction. Set to “cuda” for classifiers to be able to leverage the GPU.

  • verbose (bool, default=False) – Specifies whether to print verbose status updates of NLI scoring process

  • nli_model_name (str, default="microsoft/deberta-large-mnli") – Specifies which NLI model to use. Must be acceptable input to AutoTokenizer.from_pretrained() and AutoModelForSequenceClassification.from_pretrained()

  • max_length (int, default=2000) – Specifies the maximum allowed string length. Responses longer than this value will be truncated to avoid OutOfMemoryError. Inputs are additionally truncated at the token level to the model’s maximum sequence length.

  • batch_size (int, default=32) – Number of premise-hypothesis pairs scored per forward pass in predict_batch. Lower this value if inference runs out of memory; raise it to increase GPU utilization.

  • use_fp16 (bool, default=False) – If True and the device is CUDA or MPS, runs inference in half precision (torch.float16) to reduce memory usage and increase throughput at a small cost to numerical precision. Ignored on CPU.

  • device_map (str, default=None) – Optional device map (e.g. “auto”) passed to AutoModelForSequenceClassification.from_pretrained() to shard or place the model across available devices. Requires the accelerate package. If provided, takes precedence over device. This option is in beta and may change in future releases.

Methods

__init__([device, verbose, nli_model_name, ...])

A class to computing NLI-based confidence scores.

get_nli_results(response1, response2)

This method computes mean NLI score and determines whether entailment exists.

get_nli_results_batch(response_pairs)

Compute mean NLI scores and entailment indicators for a list of response pairs.

predict(premise, hypothesis)

This method compute probability of contradiction on the provide inputs.

predict_batch(pairs)

Compute NLI probabilities for a list of (premise, hypothesis) pairs using batched inference.
get_nli_results(response1, response2)#

This method computes mean NLI score and determines whether entailment exists.

Return type:

Dict[str, Any]

get_nli_results_batch(response_pairs)#

Compute mean NLI scores and entailment indicators for a list of response pairs. Both directions of each pair are evaluated in a single batched forward pass.

Return type:

List[Dict[str, Any]]

Parameters:

response_pairs (list of (str, str) tuples) – Response pairs to score.

Returns:

One dictionary per input pair with keys “noncontradiction_score”, “entailment”, and “entailment_score”.

Return type:

list of dict

predict(premise, hypothesis)#

This method compute probability of contradiction on the provide inputs.

Return type:

Any

Parameters:

  • premise (str) – An input for the sequence classification DeBERTa model.

  • hypothesis (str) – An input for the sequence classification DeBERTa model.

Returns:

Probabilities computed by NLI model

Return type:

numpy.ndarray

predict_batch(pairs)#

Compute NLI probabilities for a list of (premise, hypothesis) pairs using batched inference.

Return type:

Any

Parameters:

pairs (list of (str, str) tuples) – Premise-hypothesis pairs to score. Pairs are scored in chunks of batch_size per forward pass.

Returns:

Array of shape (len(pairs), 3) containing [contradiction, neutral, entailment] probabilities, one row per input pair.

Return type:

numpy.ndarray

References