chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,31 @@
|
||||
"""
|
||||
---
|
||||
title: Sampling Techniques for Language Models
|
||||
summary: >
|
||||
A set of PyTorch implementations/tutorials of sampling techniques for language models.
|
||||
---
|
||||
|
||||
# Sampling Techniques for Language Models
|
||||
|
||||
* [Greedy Sampling](greedy.html)
|
||||
* [Temperature Sampling](temperature.html)
|
||||
* [Top-k Sampling](top_k.html)
|
||||
* [Nucleus Sampling](nucleus.html)
|
||||
|
||||
Here's an [experiment](experiment.html) that uses these sampling techniques.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class Sampler:
|
||||
"""
|
||||
### Sampler base class
|
||||
"""
|
||||
def __call__(self, logits: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
### Sample from logits
|
||||
|
||||
:param logits: are the logits of the distribution of shape `[..., n_tokens]`
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,111 @@
|
||||
"""
|
||||
---
|
||||
title: Trying out Sampling Techniques for Language Models
|
||||
summary: >
|
||||
We try out different sampling techniques for language models on HuggingFace's GPT2 model.
|
||||
---
|
||||
|
||||
# Trying out Sampling Techniques for Language Models
|
||||
|
||||
* [Greedy Sampling](greedy.html)
|
||||
* [Temperature Sampling](temperature.html)
|
||||
* [Top-k Sampling](top_k.html)
|
||||
* [Nucleus Sampling](nucleus.html)
|
||||
|
||||
This experiment uses the above sampling techniques, on HuggingFace's GPT2 model.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from labml import monit, logger, lab
|
||||
|
||||
from labml.logger import Text
|
||||
|
||||
from labml_nn.sampling import Sampler
|
||||
from labml_nn.sampling.greedy import GreedySampler
|
||||
from labml_nn.sampling.nucleus import NucleusSampler
|
||||
from labml_nn.sampling.temperature import TemperatureSampler
|
||||
from labml_nn.sampling.top_k import TopKSampler
|
||||
from transformers import GPT2Tokenizer, GPT2LMHeadModel
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sample(model: GPT2LMHeadModel, tokenizer: GPT2Tokenizer, sampler: Sampler,
|
||||
n_samples: int, n_tokens: int, seq_len: int, prompt: str):
|
||||
"""
|
||||
## Sample from model
|
||||
|
||||
:param model: is the model to sample from
|
||||
:param tokenizer: is the tokenizer to use
|
||||
:param sampler: is the sampler to use
|
||||
:param n_samples: is the number of samples to generate
|
||||
:param n_tokens: is the number of tokens to generate
|
||||
:param seq_len: is the maximum sequence length for the model
|
||||
:param prompt: is the starting prompt
|
||||
"""
|
||||
# Tokenize the `prompt` and make `n_samples` copies of it
|
||||
data = torch.tile(torch.tensor(tokenizer.encode(prompt))[None, :], (n_samples, 1))
|
||||
|
||||
# Collect output for printing
|
||||
logs = [[(prompt, Text.meta)] for _ in range(n_samples)]
|
||||
# Sample `n_tokens`
|
||||
for i in monit.iterate('Sample', n_tokens):
|
||||
# Truncate the data to the maximum sequence length
|
||||
data = data[-seq_len:]
|
||||
# Get the model output. The 'logits' has shape `[batch_size, seq_len, n_tokens]`
|
||||
logits = model(data)[0]
|
||||
# Get the `logits` of the last token
|
||||
logits = logits[:, -1]
|
||||
# Sample from the `logits`
|
||||
res = sampler(logits)
|
||||
# Add the sampled token to the data
|
||||
data = torch.cat([data, res[:, None]], dim=1)
|
||||
# Decode and add the sampled token for logging
|
||||
for j in range(n_samples):
|
||||
logs[j] += [('' + tokenizer.decode(res[j]), Text.value)]
|
||||
|
||||
# Print the sampled outputs
|
||||
for j in range(n_samples):
|
||||
logger.log(logs[j])
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
### Try different sampling techniques
|
||||
"""
|
||||
|
||||
# Load the model and tokenizer
|
||||
with monit.section('Load tokenizer/model'):
|
||||
tokenizer = GPT2Tokenizer.from_pretrained('gpt2', cache_dir=lab.get_data_path() / 'cache')
|
||||
model = GPT2LMHeadModel.from_pretrained('gpt2', cache_dir=lab.get_data_path() / 'cache')
|
||||
# Set the model to eval mode
|
||||
model.eval()
|
||||
|
||||
# Prompts to use for sampling
|
||||
prompt = 'I saw an interesting dream last night. '
|
||||
|
||||
# [Greedy Sampling](greedy.html)
|
||||
with monit.section('greedy'):
|
||||
sample(model, tokenizer, GreedySampler(), 4, 32, 128, prompt)
|
||||
|
||||
# [Temperature Sampling](temperature.html)
|
||||
with monit.section('temperature=1.'):
|
||||
sample(model, tokenizer, TemperatureSampler(1.), 4, 32, 128, prompt)
|
||||
with monit.section('temperature=.1'):
|
||||
sample(model, tokenizer, TemperatureSampler(.1), 4, 32, 128, prompt)
|
||||
with monit.section('temperature=10.'):
|
||||
sample(model, tokenizer, TemperatureSampler(10.), 4, 32, 128, prompt)
|
||||
|
||||
# [Top-k Sampling](top_k.html)
|
||||
with monit.section('top_k=5'):
|
||||
sample(model, tokenizer, TopKSampler(2, TemperatureSampler(1.)), 4, 32, 128, prompt)
|
||||
|
||||
# [Nucleus Sampling](nucleus.html)
|
||||
with monit.section('nucleus p=.95'):
|
||||
sample(model, tokenizer, NucleusSampler(0.95, TemperatureSampler(1.)), 4, 32, 128, prompt)
|
||||
with monit.section('nucleus p=.1'):
|
||||
sample(model, tokenizer, NucleusSampler(0.1, TemperatureSampler(1.)), 4, 32, 128, prompt)
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,82 @@
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from labml import experiment, monit
|
||||
from labml import logger
|
||||
from labml.logger import Text
|
||||
from labml_nn.helpers.datasets import TextDataset
|
||||
from labml_nn.sampling import Sampler
|
||||
from labml_nn.sampling.greedy import GreedySampler
|
||||
from labml_nn.sampling.nucleus import NucleusSampler
|
||||
from labml_nn.sampling.temperature import TemperatureSampler
|
||||
from labml_nn.sampling.top_k import TopKSampler
|
||||
from labml_nn.transformers.basic.autoregressive_experiment import Configs, AutoregressiveTransformer
|
||||
|
||||
|
||||
def get_model_dataset(run_uuid: str) -> Tuple[AutoregressiveTransformer, TextDataset]:
|
||||
experiment.evaluate()
|
||||
|
||||
conf = Configs()
|
||||
|
||||
experiment.configs(conf, experiment.load_configs(run_uuid))
|
||||
|
||||
experiment.load(run_uuid)
|
||||
|
||||
experiment.add_pytorch_models({'model': conf.model})
|
||||
|
||||
experiment.start()
|
||||
|
||||
return conf.model, conf.text
|
||||
|
||||
|
||||
def sample(model, ds, sampler: Sampler, n_samples: int, n_tokens: int, seq_len: int, prompt: str):
|
||||
with torch.no_grad():
|
||||
data = torch.tile(ds.text_to_i(prompt)[:, None], (1, n_samples))
|
||||
|
||||
# Collect output for printing
|
||||
logs = [[(prompt, Text.meta)] for _ in range(n_samples)]
|
||||
# Sample 25 tokens
|
||||
for i in monit.iterate('Sample', n_tokens):
|
||||
# Tokenize the prompt
|
||||
data = data[-seq_len:]
|
||||
# Get the model output
|
||||
logits, *_ = model(data)
|
||||
logits = logits[-1]
|
||||
# Get the model prediction (greedy)
|
||||
res = sampler(logits)
|
||||
data = torch.cat([data, res[None, :]], dim=0)
|
||||
# Add the prediction for logging
|
||||
for j in range(n_samples):
|
||||
logs[j] += [('' + ds.itos[res[j]], Text.value)]
|
||||
|
||||
# Print the sampled output
|
||||
for j in range(n_samples):
|
||||
logger.log(logs[j])
|
||||
|
||||
|
||||
def main():
|
||||
model, ds = get_model_dataset('074d4004cc6b11ecad7a0242ac1c0002')
|
||||
model.eval()
|
||||
|
||||
with monit.section('greedy'):
|
||||
sample(model, ds, GreedySampler(), 4, 32, 128, 'It is')
|
||||
|
||||
with monit.section('temperature=1.'):
|
||||
sample(model, ds, TemperatureSampler(1.), 4, 32, 128, 'It is')
|
||||
with monit.section('temperature=.1'):
|
||||
sample(model, ds, TemperatureSampler(.1), 4, 32, 128, 'It is')
|
||||
with monit.section('temperature=10.'):
|
||||
sample(model, ds, TemperatureSampler(10.), 4, 32, 128, 'It is')
|
||||
|
||||
with monit.section('top_k=5'):
|
||||
sample(model, ds, TopKSampler(2, TemperatureSampler(1.)), 4, 32, 128, 'It is')
|
||||
|
||||
with monit.section('nucles p=.95'):
|
||||
sample(model, ds, NucleusSampler(0.95, TemperatureSampler(1.)), 4, 32, 128, 'It is')
|
||||
with monit.section('nucles p=.95'):
|
||||
sample(model, ds, NucleusSampler(0.1, TemperatureSampler(1.)), 4, 32, 128, 'It is')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,24 @@
|
||||
"""
|
||||
---
|
||||
title: Greedy Sampling
|
||||
summary: A PyTorch implementation of greedy sampling from language models.
|
||||
---
|
||||
|
||||
# Greedy Sampling
|
||||
|
||||
Here we sample the most likely token from the distribution of logits.
|
||||
|
||||
Here's an [experiment](experiment.html) that uses these sampling techniques.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from labml_nn.sampling import Sampler
|
||||
|
||||
|
||||
class GreedySampler(Sampler):
|
||||
def __call__(self, logits: torch.Tensor):
|
||||
"""
|
||||
Sample the most likely token from the distribution of logits
|
||||
"""
|
||||
return logits.argmax(dim=-1)
|
||||
@@ -0,0 +1,78 @@
|
||||
"""
|
||||
---
|
||||
title: Nucleus Sampling
|
||||
summary: A PyTorch implementation of nucleus sampling from language models.
|
||||
---
|
||||
|
||||
# Nucleus Sampling
|
||||
|
||||
This is an implementation of nucleus sampling, introduced in the paper
|
||||
[The Curious Case of Neural Text Degeneration](https://arxiv.org/abs/1904.09751).
|
||||
|
||||
The paper discusses the problems with other sampling methods such as Beam Search,
|
||||
[Pure sampling](temperature.html), [Temperature sampling](temperature.html), and
|
||||
[Top-k sampling](top_k.html). The paper introduces the idea of nucleus sampling,
|
||||
which practically performs better than other sampling methods for text generation.
|
||||
|
||||
Nucleus sampling first picks a subset of the vocabulary $V^{(p)} \subset V$,
|
||||
where $V^{(p)}$ is smallest set of tokens such that
|
||||
|
||||
$$\sum_{x_i \in V^{(p)}} P(x_i | x_{1:i-1}) \ge p$$
|
||||
|
||||
That is, we pick the highest probable tokens until the sum of their probabilities is less that $p$.
|
||||
|
||||
Then we sample from the selected tokens.
|
||||
|
||||
Here's an [experiment](experiment.html) that uses these sampling techniques.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from labml_nn.sampling import Sampler
|
||||
|
||||
|
||||
class NucleusSampler(Sampler):
|
||||
"""
|
||||
## Nucleus Sampler
|
||||
"""
|
||||
def __init__(self, p: float, sampler: Sampler):
|
||||
"""
|
||||
:param p: is the sum of probabilities of tokens to pick $p$
|
||||
:param sampler: is the sampler to use for the selected tokens
|
||||
"""
|
||||
self.p = p
|
||||
self.sampler = sampler
|
||||
# Softmax to compute $P(x_i | x_{1:i-1})$ from the logits
|
||||
self.softmax = nn.Softmax(dim=-1)
|
||||
|
||||
def __call__(self, logits: torch.Tensor):
|
||||
"""
|
||||
Sample from logits with Nucleus Sampling
|
||||
"""
|
||||
|
||||
# Get probabilities $P(x_i | x_{1:i-1})$
|
||||
probs = self.softmax(logits)
|
||||
|
||||
# Sort probabilities in descending order
|
||||
sorted_probs, indices = torch.sort(probs, dim=-1, descending=True)
|
||||
# Get the cumulative sum of probabilities in the sorted order
|
||||
cum_sum_probs = torch.cumsum(sorted_probs, dim=-1)
|
||||
# Find the cumulative sums less than $p$.
|
||||
nucleus = cum_sum_probs < self.p
|
||||
# Prepend ones so that we add one token after the minimum number
|
||||
# of tokens with cumulative probability less that $p$.
|
||||
nucleus = torch.cat([nucleus.new_ones(nucleus.shape[:-1] + (1,)), nucleus[..., :-1]], dim=-1)
|
||||
|
||||
# Get log probabilities and mask out the non-nucleus
|
||||
sorted_log_probs = torch.log(sorted_probs)
|
||||
sorted_log_probs[~nucleus] = float('-inf')
|
||||
|
||||
# Sample from the sampler
|
||||
sampled_sorted_indexes = self.sampler(sorted_log_probs)
|
||||
|
||||
# Get the actual indexes
|
||||
res = indices.gather(-1, sampled_sorted_indexes.unsqueeze(-1))
|
||||
|
||||
#
|
||||
return res.squeeze(-1)
|
||||
@@ -0,0 +1,44 @@
|
||||
"""
|
||||
---
|
||||
title: Sampling from Language Models with Temperature
|
||||
summary: A PyTorch implementation of sampling from language models with temperature.
|
||||
---
|
||||
|
||||
# Sampling from Language Models with Temperature
|
||||
|
||||
Here we sample from the following probability distribution where $V$ is the vocabulary,
|
||||
$u_{1:|V|}$ are the logits of the distribution and T is the temperature:
|
||||
|
||||
$$P(x_i=V_l | x_{1:i-1}) = \frac{\exp(\frac{u_l}{T})}{\sum_j \exp(\frac{u_j}{T})}$$
|
||||
|
||||
$T = 1$ is normal random sampling.
|
||||
|
||||
Here's an [experiment](experiment.html) that uses these sampling techniques.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch.distributions import Categorical
|
||||
|
||||
from labml_nn.sampling import Sampler
|
||||
|
||||
|
||||
class TemperatureSampler(Sampler):
|
||||
"""
|
||||
## Sampler with Temperature
|
||||
"""
|
||||
def __init__(self, temperature: float = 1.0):
|
||||
"""
|
||||
:param temperature: is the temperature to sample with
|
||||
"""
|
||||
self.temperature = temperature
|
||||
|
||||
def __call__(self, logits: torch.Tensor):
|
||||
"""
|
||||
Sample from logits
|
||||
"""
|
||||
|
||||
# Create a categorical distribution with temperature adjusted logits
|
||||
dist = Categorical(logits=logits / self.temperature)
|
||||
|
||||
# Sample
|
||||
return dist.sample()
|
||||
@@ -0,0 +1,48 @@
|
||||
"""
|
||||
---
|
||||
title: Top-k Sampling
|
||||
summary: A PyTorch implementation of top-k sampling from language models.
|
||||
---
|
||||
|
||||
# Top-k Sampling
|
||||
|
||||
Here we first pick the top-k tokens from the distribution of logits, and then
|
||||
sample from them.
|
||||
|
||||
Here's an [experiment](experiment.html) that uses these sampling techniques.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from labml_nn.sampling import Sampler
|
||||
|
||||
|
||||
class TopKSampler(Sampler):
|
||||
"""
|
||||
## Top-k Sampler
|
||||
"""
|
||||
def __init__(self, k: int, sampler: Sampler):
|
||||
"""
|
||||
:param k: is the number of tokens to pick
|
||||
:param sampler: is the sampler to use for the top-k tokens
|
||||
|
||||
`sampler` can be any sampler that takes a logits tensor as input and returns a token tensor;
|
||||
e.g. [`TemperatureSampler'](temperature.html).
|
||||
"""
|
||||
self.k = k
|
||||
self.sampler = sampler
|
||||
|
||||
def __call__(self, logits: torch.Tensor):
|
||||
"""
|
||||
Sample from logits
|
||||
"""
|
||||
# New logits filled with $-\infty$; i.e. zero probability
|
||||
zeros = logits.new_ones(logits.shape) * float('-inf')
|
||||
# Pick the largest $k$ logits and their indices
|
||||
values, indices = torch.topk(logits, self.k, dim=-1)
|
||||
# Set the values of the top-k selected indices to actual logits.
|
||||
# Logits of other tokens remain $-\infty$
|
||||
zeros.scatter_(-1, indices, values)
|
||||
|
||||
# Sample from the top-k logits with the specified sampler.
|
||||
return self.sampler(zeros)
|
||||
Reference in New Issue
Block a user