#

使用 GPT-NEOX 生成文本

这说明了如何使用单个 GPU 从 GPT-NEOX 生成文本。

这需要一个内存超过45GB的GPU。

#

进口

16from typing import List
17
18import torch
19from torch import nn
20
21from labml import monit
22from labml_nn.neox.model import LayerGenerator
23from labml_nn.neox.utils import get_tokens, print_tokens
24from labml_nn.neox.utils.cache import get_cache

#

要加载的图层列表。这用于测试。您可以将层的子集分配给变压器层，{0, 1} 使其仅将第一个层加载到变压器层。

29LAYERS = None

#

提示完成

32PROMPT = 'Einstein was born in the German Empire, but moved to Switzerland in 1895, forsaking his German'

#

预测下一个代币

model 是模特吗
ids 是输入令牌 ID
device 是该型号的设备

35def infer(model: nn.Module, ids: List[int], device: torch.device):

#

44    with torch.no_grad():

#

获取代币

46        x = torch.tensor(ids)[None, :].to(device)

#

评估模型

48        x = model(x)

#

返回预测的代币

51    return x[0].max(dim=-1)[1].tolist()

#

生成文本

54def generate():

#

设置缓存以缓存中间键/值对以加快生成速度

60    cache = get_cache()
61    cache.set('use_cache', True)

#

设备

64    device = torch.device('cuda:0')

#

加载图层

67    layers = list(LayerGenerator(is_clone_layers=True,
68                                 filter_layers=LAYERS,
69                                 dtype=torch.float16,
70                                 device=device,
71                                 ).load())
72
73    model = nn.Sequential(*layers)

#

获取代币 ID

76    ids = get_tokens(PROMPT)

#

运行模型

79    cache.set('state_ids', (None, 1))
80    with monit.section('Infer'):
81        next_token = infer(model, ids, device)[-1]

#

追加预测的令牌

84    ids += [next_token]

#

预测 100 个代币

87    for i in range(1, 100):

#

设置状态以使用缓存的激活

89        cache.set('state_ids', (i, i + 1))

#

获取下一个令牌。请注意，我们只将最后一个令牌提供给模型，因为我们缓存了先前令牌的键/值对。

92        with monit.section('Infer'):
93            next_token = infer(model, [next_token], device)[-1]

#

追加预测的令牌

95        ids += [next_token]

#

打印

97        print_tokens(ids, [ids])

#

101if __name__ == '__main__':
102    generate()