lm_polygraph.utils.causal_lm_with_uncertainty module

class lm_polygraph.utils.causal_lm_with_uncertainty.CausalLMWithUncertainty(llm, tokenizer, stat_calculators, estimator)[source]

Bases: object

device()[source]

generate(input_ids, attention_mask=None, **kwargs)[source]

class lm_polygraph.utils.causal_lm_with_uncertainty.GenerateDecoderOnlyOutputWithUncertainty(sequences: LongTensor, scores: tuple[torch.FloatTensor] | None = None, logits: tuple[torch.FloatTensor] | None = None, attentions: tuple[tuple[torch.FloatTensor]] | None = None, hidden_states: tuple[tuple[torch.FloatTensor]] | None = None, past_key_values: Cache | None = None, uncertainty_score: float | List[float] | Tensor | None = None)[source]

Bases: GenerateDecoderOnlyOutput

Extends GenerateDecoderOnlyOutput to include uncertainty scores

uncertainty_score: float | List[float] | Tensor | None = None