Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
|
@@ -1,32 +1,94 @@
|
|
| 1 |
-
import
|
| 2 |
-
import
|
| 3 |
-
|
| 4 |
-
|
| 5 |
-
#
|
| 6 |
-
#
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
def
|
| 15 |
-
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
|
| 30 |
-
|
| 31 |
-
|
| 32 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
import torch.nn as nn
|
| 3 |
+
from torch.nn import functional as F
|
| 4 |
+
|
| 5 |
+
# Параметры нашей "Альфа" версии
|
| 6 |
+
batch_size = 32 # Сколько кусочков текста учим за раз
|
| 7 |
+
block_size = 64 # Длина "памяти" (контекста) в символах
|
| 8 |
+
n_embd = 128 # Размер внутреннего "вектора мысли"
|
| 9 |
+
n_head = 4 # Количество "голов" внимания (как 4 разных взгляда на текст)
|
| 10 |
+
n_layer = 4 # Сколько слоев нейронов в глубину
|
| 11 |
+
|
| 12 |
+
class Head(nn.Module):
|
| 13 |
+
""" Одиночная голова самовнимания """
|
| 14 |
+
def __init__(self, head_size):
|
| 15 |
+
super().__init__()
|
| 16 |
+
self.key = nn.Linear(n_embd, head_size, bias=False)
|
| 17 |
+
self.query = nn.Linear(n_embd, head_size, bias=False)
|
| 18 |
+
self.value = nn.Linear(n_embd, head_size, bias=False)
|
| 19 |
+
self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
|
| 20 |
+
|
| 21 |
+
def forward(self, x):
|
| 22 |
+
B, T, C = x.shape
|
| 23 |
+
k = self.key(x)
|
| 24 |
+
q = self.query(x)
|
| 25 |
+
# Вычисляем веса внимания (на что ИИ смотрит сейчас)
|
| 26 |
+
wei = q @ k.transpose(-2,-1) * C**-0.5
|
| 27 |
+
wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
|
| 28 |
+
wei = F.softmax(wei, dim=-1)
|
| 29 |
+
# Применяем внимание к данным
|
| 30 |
+
v = self.value(x)
|
| 31 |
+
out = wei @ v
|
| 32 |
+
return out
|
| 33 |
+
|
| 34 |
+
class MultiHeadAttention(nn.Module):
|
| 35 |
+
""" Несколько голов внимания, работающих параллельно """
|
| 36 |
+
def __init__(self, num_heads, head_size):
|
| 37 |
+
super().__init__()
|
| 38 |
+
self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
|
| 39 |
+
self.proj = nn.Linear(n_embd, n_embd)
|
| 40 |
+
|
| 41 |
+
def forward(self, x):
|
| 42 |
+
out = torch.cat([h(x) for h in self.heads], dim=-1)
|
| 43 |
+
out = self.proj(out)
|
| 44 |
+
return out
|
| 45 |
+
|
| 46 |
+
class Block(nn.Module):
|
| 47 |
+
""" Один блок Трансформера: внимание + раздумья """
|
| 48 |
+
def __init__(self, n_embd, n_head):
|
| 49 |
+
super().__init__()
|
| 50 |
+
head_size = n_embd // n_head
|
| 51 |
+
self.sa = MultiHeadAttention(n_head, head_size)
|
| 52 |
+
self.ffwd = nn.Sequential(
|
| 53 |
+
nn.Linear(n_embd, 4 * n_embd),
|
| 54 |
+
nn.ReLU(),
|
| 55 |
+
nn.Linear(4 * n_embd, n_embd),
|
| 56 |
+
)
|
| 57 |
+
self.ln1 = nn.LayerNorm(n_embd)
|
| 58 |
+
self.ln2 = nn.LayerNorm(n_embd)
|
| 59 |
+
|
| 60 |
+
def forward(self, x):
|
| 61 |
+
x = x + self.sa(self.ln1(x))
|
| 62 |
+
x = x + self.ffwd(self.ln2(x))
|
| 63 |
+
return x
|
| 64 |
+
|
| 65 |
+
class AxisModel(nn.Module):
|
| 66 |
+
def __init__(self, vocab_size):
|
| 67 |
+
super().__init__()
|
| 68 |
+
# Каждому символу — свой вектор
|
| 69 |
+
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
|
| 70 |
+
# Каждой позиции в тексте — свой вектор
|
| 71 |
+
self.position_embedding_table = nn.Embedding(block_size, n_embd)
|
| 72 |
+
# Слои блоков Трансформера
|
| 73 |
+
self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)])
|
| 74 |
+
self.ln_f = nn.LayerNorm(n_embd)
|
| 75 |
+
self.lm_head = nn.Linear(n_embd, vocab_size)
|
| 76 |
+
|
| 77 |
+
def forward(self, idx, targets=None):
|
| 78 |
+
B, T = idx.shape
|
| 79 |
+
tok_emb = self.token_embedding_table(idx) # (B,T,C)
|
| 80 |
+
pos_emb = self.position_embedding_table(torch.arange(T)) # (T,C)
|
| 81 |
+
x = tok_emb + pos_emb # Объединяем смысл и позицию буквы
|
| 82 |
+
x = self.blocks(x)
|
| 83 |
+
x = self.ln_f(x)
|
| 84 |
+
logits = self.lm_head(x) # (B,T,vocab_size)
|
| 85 |
+
|
| 86 |
+
if targets is None:
|
| 87 |
+
loss = None
|
| 88 |
+
else:
|
| 89 |
+
B, T, C = logits.shape
|
| 90 |
+
logits = logits.view(B*T, C)
|
| 91 |
+
targets = targets.view(B*T)
|
| 92 |
+
loss = F.cross_entropy(logits, targets)
|
| 93 |
+
|
| 94 |
+
return logits, loss
|