5. Архітектура LLM

Reading time: 16 minutes

Архітектура LLM

Приклад архітектури LLM з https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb:

Високорівневе подання можна спостерігати в:

1. Вхід (Токенізований текст): Процес починається з токенізованого тексту, який перетворюється на числові представлення.
2. Шар вбудовування токенів та шар позиційного вбудовування: Токенізований текст проходить через шар вбудовування токенів та шар позиційного вбудовування, який захоплює позицію токенів у послідовності, що є критично важливим для розуміння порядку слів.
3. Блоки трансформера: Модель містить 12 блоків трансформера, кожен з яких має кілька шарів. Ці блоки повторюють наступну послідовність:

• Масковане багатоголове увага: Дозволяє моделі зосереджуватися на різних частинах вхідного тексту одночасно.
• Нормалізація шару: Крок нормалізації для стабілізації та покращення навчання.
• Шар прямого проходження: Відповідає за обробку інформації з шару уваги та прогнозування наступного токена.
• Шари відсіву: Ці шари запобігають перенавчанню, випадковим чином відкидаючи одиниці під час навчання.

4. Фінальний вихідний шар: Модель виводить тензор розміром 4x50,257, де 50,257 представляє розмір словника. Кожен рядок у цьому тензорі відповідає вектору, який модель використовує для прогнозування наступного слова в послідовності.
5. Мета: Завдання полягає в тому, щоб взяти ці вбудовування та перетворити їх назад у текст. Конкретно, останній рядок виходу використовується для генерації наступного слова, представленого як "вперед" у цій діаграмі.

Подання коду

import torch
import torch.nn as nn
import tiktoken

class GELU(nn.Module):
def __init__(self):
super().__init__()

def forward(self, x):
return 0.5 * x * (1 + torch.tanh(
torch.sqrt(torch.tensor(2.0 / torch.pi)) *
(x + 0.044715 * torch.pow(x, 3))
))

class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
GELU(),
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
)

def forward(self, x):
return self.layers(x)

class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
assert d_out % num_heads == 0, "d_out must be divisible by num_heads"

self.d_out = d_out
self.num_heads = num_heads
self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim

self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
self.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
self.dropout = nn.Dropout(dropout)
self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))

def forward(self, x):
b, num_tokens, d_in = x.shape

keys = self.W_key(x) # Shape: (b, num_tokens, d_out)
queries = self.W_query(x)
values = self.W_value(x)

# We implicitly split the matrix by adding a `num_heads` dimension
# Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
values = values.view(b, num_tokens, self.num_heads, self.head_dim)
queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

# Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
keys = keys.transpose(1, 2)
queries = queries.transpose(1, 2)
values = values.transpose(1, 2)

# Compute scaled dot-product attention (aka self-attention) with a causal mask
attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head

# Original mask truncated to the number of tokens and converted to boolean
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]

# Use the mask to fill attention scores
attn_scores.masked_fill_(mask_bool, -torch.inf)

attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
attn_weights = self.dropout(attn_weights)

# Shape: (b, num_tokens, num_heads, head_dim)
context_vec = (attn_weights @ values).transpose(1, 2)

# Combine heads, where self.d_out = self.num_heads * self.head_dim
context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
context_vec = self.out_proj(context_vec) # optional projection

return context_vec

class LayerNorm(nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.eps = 1e-5
self.scale = nn.Parameter(torch.ones(emb_dim))
self.shift = nn.Parameter(torch.zeros(emb_dim))

def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift

class TransformerBlock(nn.Module):
def __init__(self, cfg):
super().__init__()
self.att = MultiHeadAttention(
d_in=cfg["emb_dim"],
d_out=cfg["emb_dim"],
context_length=cfg["context_length"],
num_heads=cfg["n_heads"],
dropout=cfg["drop_rate"],
qkv_bias=cfg["qkv_bias"])
self.ff = FeedForward(cfg)
self.norm1 = LayerNorm(cfg["emb_dim"])
self.norm2 = LayerNorm(cfg["emb_dim"])
self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

def forward(self, x):
# Shortcut connection for attention block
shortcut = x
x = self.norm1(x)
x = self.att(x)  # Shape [batch_size, num_tokens, emb_size]
x = self.drop_shortcut(x)
x = x + shortcut  # Add the original input back

# Shortcut connection for feed forward block
shortcut = x
x = self.norm2(x)
x = self.ff(x)
x = self.drop_shortcut(x)
x = x + shortcut  # Add the original input back

return x


class GPTModel(nn.Module):
def __init__(self, cfg):
super().__init__()
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
self.drop_emb = nn.Dropout(cfg["drop_rate"])

self.trf_blocks = nn.Sequential(
*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(
cfg["emb_dim"], cfg["vocab_size"], bias=False
)

def forward(self, in_idx):
batch_size, seq_len = in_idx.shape
tok_embeds = self.tok_emb(in_idx)
pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
x = self.drop_emb(x)
x = self.trf_blocks(x)
x = self.final_norm(x)
logits = self.out_head(x)
return logits

GPT_CONFIG_124M = {
"vocab_size": 50257,    # Vocabulary size
"context_length": 1024, # Context length
"emb_dim": 768,         # Embedding dimension
"n_heads": 12,          # Number of attention heads
"n_layers": 12,         # Number of layers
"drop_rate": 0.1,       # Dropout rate
"qkv_bias": False       # Query-Key-Value bias
}

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
out = model(batch)
print("Input batch:\n", batch)
print("\nOutput shape:", out.shape)
print(out)

Функція активації GELU

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class GELU(nn.Module):
def __init__(self):
super().__init__()

def forward(self, x):
return 0.5 * x * (1 + torch.tanh(
torch.sqrt(torch.tensor(2.0 / torch.pi)) *
(x + 0.044715 * torch.pow(x, 3))
))

Мета та Функціональність

• GELU (Гаусова Помилкова Лінійна Одиниця): Активаційна функція, яка вводить нелінійність у модель.
• Плавна Активація: На відміну від ReLU, яка обнуляє негативні вхідні дані, GELU плавно відображає вхідні дані на виходи, дозволяючи невеликі, ненульові значення для негативних вхідних даних.
• Математичне Визначення:

FeedForward Нейронна Мережа

Форми були додані як коментарі для кращого розуміння форм матриць:

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
GELU(),
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
)

def forward(self, x):
# x shape: (batch_size, seq_len, emb_dim)

x = self.layers[0](x)# x shape: (batch_size, seq_len, 4 * emb_dim)
x = self.layers[1](x) # x shape remains: (batch_size, seq_len, 4 * emb_dim)
x = self.layers[2](x) # x shape: (batch_size, seq_len, emb_dim)
return x  # Output shape: (batch_size, seq_len, emb_dim)

Мета та Функціональність

• Позиційна FeedForward мережа: Застосовує двошарову повністю з'єднану мережу до кожної позиції окремо та ідентично.
• Деталі Шарів:
• Перший Лінійний Шар: Розширює розмірність з emb_dim до 4 * emb_dim.
• Активація GELU: Застосовує нелінійність.
• Другий Лінійний Шар: Зменшує розмірність назад до emb_dim.

Механізм Багатоголової Уваги

Це вже було пояснено в попередньому розділі.

Мета та Функціональність

• Багатоголова Самоувага: Дозволяє моделі зосереджуватися на різних позиціях у вхідній послідовності під час кодування токена.
• Ключові Компоненти:
• Запити, Ключі, Значення: Лінійні проекції вхідних даних, які використовуються для обчислення оцінок уваги.
• Голови: Кілька механізмів уваги, що працюють паралельно (num_heads), кожен з зменшеною розмірністю (head_dim).
• Оцінки Уваги: Обчислюються як скалярний добуток запитів і ключів, масштабованих і замаскованих.
• Маскування: Застосовується каузальна маска, щоб запобігти моделі звертатися до майбутніх токенів (важливо для авторегресивних моделей, таких як GPT).
• Ваги Уваги: Softmax замаскованих і масштабованих оцінок уваги.
• Контекстний Вектор: Вагова сума значень відповідно до ваг уваги.
• Вихідна Проекція: Лінійний шар для об'єднання виходів усіх голів.

Нормалізація Шарів

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class LayerNorm(nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.eps = 1e-5 # Prevent division by zero during normalization.
self.scale = nn.Parameter(torch.ones(emb_dim))
self.shift = nn.Parameter(torch.zeros(emb_dim))

def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift

Мета та Функціональність

• Layer Normalization: Техніка, що використовується для нормалізації вхідних даних по ознаках (виміри вбудовування) для кожного окремого прикладу в партії.
• Компоненти:
• eps: Маленька константа (1e-5), що додається до дисперсії, щоб запобігти діленню на нуль під час нормалізації.
• scale та shift: Навчальні параметри (nn.Parameter), які дозволяють моделі масштабувати та зміщувати нормалізований вихід. Вони ініціалізуються одиницями та нулями відповідно.
• Процес Нормалізації:
• Обчислення Середнього (mean): Обчислює середнє значення вхідного x по виміру вбудовування (dim=-1), зберігаючи вимір для трансляції (keepdim=True).
• Обчислення Дисперсії (var): Обчислює дисперсію x по виміру вбудовування, також зберігаючи вимір. Параметр unbiased=False забезпечує обчислення дисперсії за допомогою упередженого оцінювача (ділення на N замість N-1), що є доречним при нормалізації по ознаках, а не зразках.
• Нормалізація (norm_x): Від x віднімається середнє значення і ділиться на квадратний корінь з дисперсії плюс eps.
• Масштабування та Зміщення: Застосовує навчальні параметри scale та shift до нормалізованого виходу.

Блок Трансформера

Форми були додані як коментарі для кращого розуміння форм матриць:

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04

class TransformerBlock(nn.Module):
def __init__(self, cfg):
super().__init__()
self.att = MultiHeadAttention(
d_in=cfg["emb_dim"],
d_out=cfg["emb_dim"],
context_length=cfg["context_length"],
num_heads=cfg["n_heads"],
dropout=cfg["drop_rate"],
qkv_bias=cfg["qkv_bias"]
)
self.ff = FeedForward(cfg)
self.norm1 = LayerNorm(cfg["emb_dim"])
self.norm2 = LayerNorm(cfg["emb_dim"])
self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

def forward(self, x):
# x shape: (batch_size, seq_len, emb_dim)

# Shortcut connection for attention block
shortcut = x  # shape: (batch_size, seq_len, emb_dim)
x = self.norm1(x)  # shape remains (batch_size, seq_len, emb_dim)
x = self.att(x)    # shape: (batch_size, seq_len, emb_dim)
x = self.drop_shortcut(x)  # shape remains (batch_size, seq_len, emb_dim)
x = x + shortcut   # shape: (batch_size, seq_len, emb_dim)

# Shortcut connection for feedforward block
shortcut = x       # shape: (batch_size, seq_len, emb_dim)
x = self.norm2(x)  # shape remains (batch_size, seq_len, emb_dim)
x = self.ff(x)     # shape: (batch_size, seq_len, emb_dim)
x = self.drop_shortcut(x)  # shape remains (batch_size, seq_len, emb_dim)
x = x + shortcut   # shape: (batch_size, seq_len, emb_dim)

return x  # Output shape: (batch_size, seq_len, emb_dim)

Мета та Функціональність

• Складові Шарів: Поєднує багатоголову увагу, мережу прямого проходження, нормалізацію шару та залишкові з'єднання.
• Нормалізація Шару: Застосовується перед шарами уваги та прямого проходження для стабільного навчання.
• Залишкові З'єднання (Швидкі З'єднання): Додають вхід шару до його виходу для покращення потоку градієнтів та можливості навчання глибоких мереж.
• Випадкове Випадання: Застосовується після шарів уваги та прямого проходження для регуляризації.

Покрокова Функціональність

1. Перший Залишковий Шлях (Само-Увага):

• Вхід (shortcut): Зберегти оригінальний вхід для залишкового з'єднання.
• Нормалізація Шару (norm1): Нормалізувати вхід.
• Багатоголова Увага (att): Застосувати само-увагу.
• Випадкове Випадання (drop_shortcut): Застосувати випадкове випадання для регуляризації.
• Додати Залишок (x + shortcut): Поєднати з оригінальним входом.

2. Другий Залишковий Шлях (Пряме Проходження):

• Вхід (shortcut): Зберегти оновлений вхід для наступного залишкового з'єднання.
• Нормалізація Шару (norm2): Нормалізувати вхід.
• Мережа Прямого Проходження (ff): Застосувати трансформацію прямого проходження.
• Випадкове Випадання (drop_shortcut): Застосувати випадкове випадання.
• Додати Залишок (x + shortcut): Поєднати з входом з першого залишкового шляху.

GPTModel

Форми були додані як коментарі для кращого розуміння форм матриць:

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class GPTModel(nn.Module):
def __init__(self, cfg):
super().__init__()
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
# shape: (vocab_size, emb_dim)

self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
# shape: (context_length, emb_dim)

self.drop_emb = nn.Dropout(cfg["drop_rate"])

self.trf_blocks = nn.Sequential(
*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
)
# Stack of TransformerBlocks

self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
# shape: (emb_dim, vocab_size)

def forward(self, in_idx):
# in_idx shape: (batch_size, seq_len)
batch_size, seq_len = in_idx.shape

# Token embeddings
tok_embeds = self.tok_emb(in_idx)
# shape: (batch_size, seq_len, emb_dim)

# Positional embeddings
pos_indices = torch.arange(seq_len, device=in_idx.device)
# shape: (seq_len,)
pos_embeds = self.pos_emb(pos_indices)
# shape: (seq_len, emb_dim)

# Add token and positional embeddings
x = tok_embeds + pos_embeds  # Broadcasting over batch dimension
# x shape: (batch_size, seq_len, emb_dim)

x = self.drop_emb(x)  # Dropout applied
# x shape remains: (batch_size, seq_len, emb_dim)

x = self.trf_blocks(x)  # Pass through Transformer blocks
# x shape remains: (batch_size, seq_len, emb_dim)

x = self.final_norm(x)  # Final LayerNorm
# x shape remains: (batch_size, seq_len, emb_dim)

logits = self.out_head(x)  # Project to vocabulary size
# logits shape: (batch_size, seq_len, vocab_size)

return logits  # Output shape: (batch_size, seq_len, vocab_size)

Мета та Функціональність

• Embedding Layers:
• Token Embeddings (tok_emb): Перетворює індекси токенів на embeddings. Нагадаємо, це ваги, які надаються кожному виміру кожного токена в словнику.
• Positional Embeddings (pos_emb): Додає позиційну інформацію до embeddings, щоб зафіксувати порядок токенів. Нагадаємо, це ваги, які надаються токену відповідно до його позиції в тексті.
• Dropout (drop_emb): Застосовується до embeddings для регуляризації.
• Transformer Blocks (trf_blocks): Стек з n_layers трансформерних блоків для обробки embeddings.
• Final Normalization (final_norm): Нормалізація шару перед вихідним шаром.
• Output Layer (out_head): Проектує фінальні приховані стани на розмір словника для отримання логітів для прогнозування.

Кількість параметрів для навчання

Маючи визначену структуру GPT, можна дізнатися кількість параметрів для навчання:

GPT_CONFIG_124M = {
"vocab_size": 50257,    # Vocabulary size
"context_length": 1024, # Context length
"emb_dim": 768,         # Embedding dimension
"n_heads": 12,          # Number of attention heads
"n_layers": 12,         # Number of layers
"drop_rate": 0.1,       # Dropout rate
"qkv_bias": False       # Query-Key-Value bias
}

model = GPTModel(GPT_CONFIG_124M)
total_params = sum(p.numel() for p in model.parameters())
print(f"Total number of parameters: {total_params:,}")
# Total number of parameters: 163,009,536

Покроковий розрахунок

1. Вкладкові шари: Вкладення токенів та позиційне вкладення

• Шар: nn.Embedding(vocab_size, emb_dim)
• Параметри: vocab_size * emb_dim

token_embedding_params = 50257 * 768 = 38,597,376

• Шар: nn.Embedding(context_length, emb_dim)
• Параметри: context_length * emb_dim

position_embedding_params = 1024 * 768 = 786,432

Загальна кількість параметрів вбудовування

embedding_params = token_embedding_params + position_embedding_params
embedding_params = 38,597,376 + 786,432 = 39,383,808

2. Transformer Blocks

Є 12 блоків трансформера, тому ми розрахуємо параметри для одного блоку, а потім помножимо на 12.

Параметри на один блок трансформера

a. Багатоголове увага

• Компоненти:

• Лінійний шар запиту (W_query): nn.Linear(emb_dim, emb_dim, bias=False)

• Лінійний шар ключа (W_key): nn.Linear(emb_dim, emb_dim, bias=False)

• Лінійний шар значення (W_value): nn.Linear(emb_dim, emb_dim, bias=False)

• Вихідна проекція (out_proj): nn.Linear(emb_dim, emb_dim)

• Розрахунки:

• Кожен з W_query, W_key, W_value:

qkv_params = emb_dim * emb_dim = 768 * 768 = 589,824

Оскільки є три такі шари:

total_qkv_params = 3 * qkv_params = 3 * 589,824 = 1,769,472

• Вихідна проекція (out_proj):

out_proj_params = (emb_dim * emb_dim) + emb_dim = (768 * 768) + 768 = 589,824 + 768 = 590,592

• Загальна кількість параметрів багатоголового уваги:

mha_params = total_qkv_params + out_proj_params
mha_params = 1,769,472 + 590,592 = 2,360,064

b. Мережа з прямим проходженням

• Компоненти:

• Перший лінійний шар: nn.Linear(emb_dim, 4 * emb_dim)

• Другий лінійний шар: nn.Linear(4 * emb_dim, emb_dim)

• Розрахунки:

• Перший лінійний шар:

ff_first_layer_params = (emb_dim * 4 * emb_dim) + (4 * emb_dim)
ff_first_layer_params = (768 * 3072) + 3072 = 2,359,296 + 3,072 = 2,362,368

• Другий лінійний шар:

ff_second_layer_params = (4 * emb_dim * emb_dim) + emb_dim
ff_second_layer_params = (3072 * 768) + 768 = 2,359,296 + 768 = 2,360,064

• Загальна кількість параметрів мережі з прямим проходженням:

ff_params = ff_first_layer_params + ff_second_layer_params
ff_params = 2,362,368 + 2,360,064 = 4,722,432

c. Нормалізації шару

• Компоненти:
• Два екземпляри LayerNorm на блок.
• Кожен LayerNorm має 2 * emb_dim параметрів (масштаб і зсув).
• Розрахунки:

layer_norm_params_per_block = 2 * (2 * emb_dim) = 2 * 768 * 2 = 3,072

d. Загальна кількість параметрів на один блок трансформера

pythonCopy codeparams_per_block = mha_params + ff_params + layer_norm_params_per_block
params_per_block = 2,360,064 + 4,722,432 + 3,072 = 7,085,568

Загальна кількість параметрів для всіх блоків трансформера

pythonCopy codetotal_transformer_blocks_params = params_per_block * n_layers
total_transformer_blocks_params = 7,085,568 * 12 = 85,026,816

3. Остаточні шари

a. Нормалізація остаточного шару

• Параметри: 2 * emb_dim (масштаб і зсув)

pythonCopy codefinal_layer_norm_params = 2 * 768 = 1,536

b. Шар виводу проекції (out_head)

• Шар: nn.Linear(emb_dim, vocab_size, bias=False)
• Параметри: emb_dim * vocab_size

pythonCopy codeoutput_projection_params = 768 * 50257 = 38,597,376

4. Підсумування всіх параметрів

pythonCopy codetotal_params = (
embedding_params +
total_transformer_blocks_params +
final_layer_norm_params +
output_projection_params
)
total_params = (
39,383,808 +
85,026,816 +
1,536 +
38,597,376
)
total_params = 163,009,536

Генерація тексту

Маючи модель, яка передбачає наступний токен, як і попередній, потрібно просто взяти значення останнього токена з виходу (оскільки вони будуть значеннями передбаченого токена), що буде значенням на запис у словнику, а потім використати функцію softmax, щоб нормалізувати виміри в ймовірності, які в сумі дорівнюють 1, а потім отримати індекс найбільшого запису, який буде індексом слова в словнику.

Code from https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb:

def generate_text_simple(model, idx, max_new_tokens, context_size):
# idx is (batch, n_tokens) array of indices in the current context
for _ in range(max_new_tokens):

# Crop current context if it exceeds the supported context size
# E.g., if LLM supports only 5 tokens, and the context size is 10
# then only the last 5 tokens are used as context
idx_cond = idx[:, -context_size:]

# Get the predictions
with torch.no_grad():
logits = model(idx_cond)

# Focus only on the last time step
# (batch, n_tokens, vocab_size) becomes (batch, vocab_size)
logits = logits[:, -1, :]

# Apply softmax to get probabilities
probas = torch.softmax(logits, dim=-1)  # (batch, vocab_size)

# Get the idx of the vocab entry with the highest probability value
idx_next = torch.argmax(probas, dim=-1, keepdim=True)  # (batch, 1)

# Append sampled index to the running sequence
idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)

return idx


start_context = "Hello, I am"

encoded = tokenizer.encode(start_context)
print("encoded:", encoded)

encoded_tensor = torch.tensor(encoded).unsqueeze(0)
print("encoded_tensor.shape:", encoded_tensor.shape)

model.eval() # disable dropout

out = generate_text_simple(
model=model,
idx=encoded_tensor,
max_new_tokens=6,
context_size=GPT_CONFIG_124M["context_length"]
)

print("Output:", out)
print("Output length:", len(out[0]))

Посилання

• https://www.manning.com/books/build-a-large-language-model-from-scratch