භාෂාආකෘති සඳහා නියැදි ක්රම අත්හදා බැලීම

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

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

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

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

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

54        data = data[-seq_len:]

56        logits = model(data)[0]

58        logits = logits[:, -1]

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)

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()