LLM Sampler

A utility library for sampling responses and extracting log-probabilities from HuggingFace causal language models.

Installation

pip install llm-sampler

Features

• Response Sampling: Sample diverse responses from causal language models
• Log-Probability Extraction: Extract log-probabilities for generated tokens
• HuggingFace Integration: Works seamlessly with HuggingFace transformers
• Batch Processing: Efficient batch sampling and probability computation

Usage

from llm_sampler import LMSampler

sampler = LMSampler("gpt2")
responses, log_probs = sampler.sample(prompt, num_samples=5)

Dependencies

• torch >= 2.0
• transformers >= 4.31
• sentencepiece

License

See LICENSE file for details.