HFTextGenerationSampler

class probly.representer.sampler.huggingface.HFTextGenerationSampler(model: GenerativePreTrainedModel, tokenizer: PreTrainedTokenizerBase, num_samples: int, batch_size: int | None = None, apply_chat_template: bool = True, add_generation_prompt: bool = True, strip_inputs: bool = True, do_sample: bool = True, temperature: float | None = None, max_new_tokens: int | None = None, top_k: int | None = None, generation_config: GenerationConfig | None = None, with_log_likelihood: bool = True, length_normalization: bool = True)

Bases: Representer[Any, (TextGenerationInputs,), Tensor, TextGenerationOutput]

Sample decoded text generations from a transformers generation model.

Initialize the text generation sampler.

Parameters:

• model – A transformers text generation model.

• tokenizer – The tokenizer associated with model.

• num_samples – Number of completions to sample for each input.

• batch_size – Number of samples per input to generate at once. None means num_samples.

• apply_chat_template – Whether inputs should use tokenizer chat templating. String inputs are wrapped as user messages; chat interactions are templated directly.

• add_generation_prompt – Whether to add a generation prompt when applying the chat template. Only has an effect if apply_chat_template=True.

• strip_inputs – Whether decoded generations should omit decoder-only prompt tokens.

• do_sample – Whether stochastic sampling should be enabled in model.generate.

• temperature – Optional sampling temperature.

• max_new_tokens – Optional maximum number of generated tokens.

• top_k – Optional top-k sampling cutoff.

• generation_config – Optional generation config for less common generation options.

• with_log_likelihood – Whether to compute generated-sequence log likelihoods.

• length_normalization – Whether log likelihoods should be averaged over scored tokens.

__call__(*args: In.args, **kwargs: In.kwargs) → R

Alias for the represent method.

add_generation_prompt: bool

apply_chat_template: bool

batch_size: int | None

do_sample: bool

classmethod from_model_name(model_name: str, num_samples: int, *, cache_dir: str | PathLike[str] | None = None, force_download: bool = False, model_kwargs: Mapping[str, object] | None = None, tokenizer_kwargs: Mapping[str, object] | None = None, batch_size: int | None = None, apply_chat_template: bool = True, add_generation_prompt: bool = True, strip_inputs: bool = True, do_sample: bool = True, temperature: float | None = None, max_new_tokens: int | None = None, top_k: int | None = None, generation_config: GenerationConfig | None = None, with_log_likelihood: bool = True, length_normalization: bool = True) → Self

Load a model by name and initialize a text generation sampler.

Parameters:

• model_name – Hugging Face model name or local model path.

• num_samples – Number of completions to sample for each input.

• cache_dir – Optional Hugging Face cache directory.

• force_download – Whether Hugging Face should re-download files even if cached.

• model_kwargs – Additional keyword arguments forwarded to AutoModelForCausalLM.from_pretrained.

• tokenizer_kwargs – Additional keyword arguments forwarded to AutoTokenizer.from_pretrained.

• batch_size – Number of samples per input to generate at once. None means num_samples.

• apply_chat_template – Whether inputs should use tokenizer chat templating. String inputs are wrapped as user messages; chat interactions are templated directly.

• add_generation_prompt – Whether to add a generation prompt when applying the chat template. Only has an effect if apply_chat_template=True.

• strip_inputs – Whether decoded generations should omit decoder-only prompt tokens.

• do_sample – Whether stochastic sampling should be enabled in model.generate.

• temperature – Optional sampling temperature.

• max_new_tokens – Optional maximum number of generated tokens.

• top_k – Optional top-k sampling cutoff.

• generation_config – Optional generation config for less common generation options.

• with_log_likelihood – Whether to compute generated-sequence log likelihoods.

• length_normalization – Whether log likelihoods should be averaged over scored tokens.

Returns:

A sampler backed by the loaded model and tokenizer.

generation_config: GenerationConfig | None

length_normalization: bool

max_new_tokens: int | None

model: GenerativePreTrainedModel

num_samples: int

predict(*args: In.args, **kwargs: In.kwargs) → R

Predict the representation for a given input.

property predictor: GenerativePreTrainedModel

The underlying model used for generation.

represent(inputs: TextGenerationInputs) → TextGenerationOutput

Sample completions for each input.

Parameters:

inputs – Prompt strings, chat-template compatible interactions, string arrays, or text generations.

Returns:

A sample with a newly added trailing sample axis.

strip_inputs: bool

temperature: float | None

tokenizer: PreTrainedTokenizerBase

top_k: int | None

with_log_likelihood: bool