#

尝试语言模型的采样技术

本实验在HuggingFace的GPT2模型上使用了上述采样技术。

18import torch
19
20from labml import monit, logger, lab
21
22from labml.logger import Text
23
24from labml_nn.sampling import Sampler
25from labml_nn.sampling.greedy import GreedySampler
26from labml_nn.sampling.nucleus import NucleusSampler
27from labml_nn.sampling.temperature import TemperatureSampler
28from labml_nn.sampling.top_k import TopKSampler
29from transformers import GPT2Tokenizer, GPT2LMHeadModel

#

来自模型的样本

model 是要采样的模型
tokenizer 是要使用的分词器
sampler 是要使用的采样器
n_samples 是要生成的样本数
n_tokens 是要生成的令牌数量
seq_len 是模型的最大序列长度
prompt 是起始提示

32@torch.no_grad()
33def sample(model: GPT2LMHeadModel, tokenizer: GPT2Tokenizer, sampler: Sampler,
34           n_samples: int, n_tokens: int, seq_len: int, prompt: str):

#

标记化prompt 并制作其n_samples 副本

47    data = torch.tile(torch.tensor(tokenizer.encode(prompt))[None, :], (n_samples, 1))

#

收集输出以进行打印

50    logs = [[(prompt, Text.meta)] for _ in range(n_samples)]

#

样本n_tokens

52    for i in monit.iterate('Sample', n_tokens):

#

将数据截断为最大序列长度

54        data = data[-seq_len:]

#

获取模型输出。“logits” 有形状[batch_size, seq_len, n_tokens]

56        logits = model(data)[0]

#

获取最后logits 一个令牌的

58        logits = logits[:, -1]

#

样本来自logits

60        res = sampler(logits)

#

将采样令牌添加到数据中

62        data = torch.cat([data, res[:, None]], dim=1)

#

解码并添加用于日志记录的采样令牌

64        for j in range(n_samples):
65            logs[j] += [('' + tokenizer.decode(res[j]), Text.value)]

#

打印采样输出

68    for j in range(n_samples):
69        logger.log(logs[j])

#

尝试不同的采样技巧

72def main():

#

加载模型和分词器

78    with monit.section('Load tokenizer/model'):
79        tokenizer = GPT2Tokenizer.from_pretrained('gpt2', cache_dir=lab.get_data_path() / 'cache')
80        model = GPT2LMHeadModel.from_pretrained('gpt2', cache_dir=lab.get_data_path() / 'cache')

#

将模型设置为评估模式

82    model.eval()

#

采样时使用的提示

85    prompt = 'I saw an interesting dream last night. '

#

贪婪采样

88    with monit.section('greedy'):
89        sample(model, tokenizer, GreedySampler(), 4, 32, 128, prompt)

#

温度采样

92    with monit.section('temperature=1.'):
93        sample(model, tokenizer, TemperatureSampler(1.), 4, 32, 128, prompt)
94    with monit.section('temperature=.1'):
95        sample(model, tokenizer, TemperatureSampler(.1), 4, 32, 128, prompt)
96    with monit.section('temperature=10.'):
97        sample(model, tokenizer, TemperatureSampler(10.), 4, 32, 128, prompt)

#

前 k 个采样

100    with monit.section('top_k=5'):
101        sample(model, tokenizer, TopKSampler(2, TemperatureSampler(1.)), 4, 32, 128, prompt)

#

原子核采样

104    with monit.section('nucleus p=.95'):
105        sample(model, tokenizer, NucleusSampler(0.95, TemperatureSampler(1.)), 4, 32, 128, prompt)
106    with monit.section('nucleus p=.1'):
107        sample(model, tokenizer, NucleusSampler(0.1, TemperatureSampler(1.)), 4, 32, 128, prompt)

#

110if __name__ == '__main__':
111    main()