tokencraft-gpt 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- tokencraft_gpt-0.1.0/PKG-INFO +34 -0
- tokencraft_gpt-0.1.0/README.md +25 -0
- tokencraft_gpt-0.1.0/config.py +0 -0
- tokencraft_gpt-0.1.0/data/data.txt +17 -0
- tokencraft_gpt-0.1.0/generate.py +135 -0
- tokencraft_gpt-0.1.0/model/__init__.py +5 -0
- tokencraft_gpt-0.1.0/model/embedding.py +35 -0
- tokencraft_gpt-0.1.0/model/feed_forward.py +39 -0
- tokencraft_gpt-0.1.0/model/gpt.py +86 -0
- tokencraft_gpt-0.1.0/model/input_embedding.py +54 -0
- tokencraft_gpt-0.1.0/model/multi_head_attention.py +117 -0
- tokencraft_gpt-0.1.0/model/positional_embedding.py +45 -0
- tokencraft_gpt-0.1.0/model/self_attention.py +64 -0
- tokencraft_gpt-0.1.0/model/transformer_block.py +67 -0
- tokencraft_gpt-0.1.0/pyproject.toml +17 -0
- tokencraft_gpt-0.1.0/tokencraft/__init__.py +6 -0
- tokencraft_gpt-0.1.0/tokencraft/model/__init__.py +3 -0
- tokencraft_gpt-0.1.0/tokencraft/model/embedding.py +35 -0
- tokencraft_gpt-0.1.0/tokencraft/model/feed_forward.py +39 -0
- tokencraft_gpt-0.1.0/tokencraft/model/gpt.py +86 -0
- tokencraft_gpt-0.1.0/tokencraft/model/input_embedding.py +54 -0
- tokencraft_gpt-0.1.0/tokencraft/model/multi_head_attention.py +117 -0
- tokencraft_gpt-0.1.0/tokencraft/model/positional_embedding.py +45 -0
- tokencraft_gpt-0.1.0/tokencraft/model/self_attention.py +64 -0
- tokencraft_gpt-0.1.0/tokencraft/model/transformer_block.py +67 -0
- tokencraft_gpt-0.1.0/tokencraft/tokenizer.py +55 -0
- tokencraft_gpt-0.1.0/tokencraft_model.pth +0 -0
- tokencraft_gpt-0.1.0/tokenizer.py +55 -0
- tokencraft_gpt-0.1.0/train.py +180 -0
- tokencraft_gpt-0.1.0/training/dataset.py +25 -0
- tokencraft_gpt-0.1.0/training/prepare_data.py +29 -0
- tokencraft_gpt-0.1.0/training/test_dataset.py +32 -0
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: tokencraft-gpt
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: An educational GPT language model implemented from scratch with PyTorch.
|
|
5
|
+
License-Expression: MIT
|
|
6
|
+
Requires-Python: >=3.10
|
|
7
|
+
Requires-Dist: torch>=2.0
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
|
|
10
|
+
cat > README.md <<'EOF'
|
|
11
|
+
# TokenCraft
|
|
12
|
+
|
|
13
|
+
**From tokens to intelligence.**
|
|
14
|
+
|
|
15
|
+
TokenCraft is a small educational GPT-style language model implemented using PyTorch.
|
|
16
|
+
|
|
17
|
+
The project demonstrates how a GPT model works, from tokenization to next-token generation.
|
|
18
|
+
|
|
19
|
+
## Features
|
|
20
|
+
|
|
21
|
+
- Simple word-level tokenizer
|
|
22
|
+
- Token and positional embeddings
|
|
23
|
+
- Causal self-attention
|
|
24
|
+
- Multi-head attention
|
|
25
|
+
- Feed-forward network
|
|
26
|
+
- Transformer blocks
|
|
27
|
+
- GPT decoder architecture
|
|
28
|
+
- Training with next-token prediction
|
|
29
|
+
- Text generation with temperature sampling
|
|
30
|
+
|
|
31
|
+
## Installation
|
|
32
|
+
|
|
33
|
+
```bash
|
|
34
|
+
pip install tokencraft-gpt
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
cat > README.md <<'EOF'
|
|
2
|
+
# TokenCraft
|
|
3
|
+
|
|
4
|
+
**From tokens to intelligence.**
|
|
5
|
+
|
|
6
|
+
TokenCraft is a small educational GPT-style language model implemented using PyTorch.
|
|
7
|
+
|
|
8
|
+
The project demonstrates how a GPT model works, from tokenization to next-token generation.
|
|
9
|
+
|
|
10
|
+
## Features
|
|
11
|
+
|
|
12
|
+
- Simple word-level tokenizer
|
|
13
|
+
- Token and positional embeddings
|
|
14
|
+
- Causal self-attention
|
|
15
|
+
- Multi-head attention
|
|
16
|
+
- Feed-forward network
|
|
17
|
+
- Transformer blocks
|
|
18
|
+
- GPT decoder architecture
|
|
19
|
+
- Training with next-token prediction
|
|
20
|
+
- Text generation with temperature sampling
|
|
21
|
+
|
|
22
|
+
## Installation
|
|
23
|
+
|
|
24
|
+
```bash
|
|
25
|
+
pip install tokencraft-gpt
|
|
File without changes
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
A language model learns patterns from text.
|
|
2
|
+
A language model predicts the next token.
|
|
3
|
+
A token can be a word or a piece of a word.
|
|
4
|
+
The tokenizer converts text into token IDs.
|
|
5
|
+
The model converts token IDs into embeddings.
|
|
6
|
+
Embeddings represent tokens as vectors.
|
|
7
|
+
A transformer uses attention to process context.
|
|
8
|
+
Self attention helps the model understand relationships between tokens.
|
|
9
|
+
Causal attention prevents the model from seeing future tokens.
|
|
10
|
+
The feed forward network processes token representations.
|
|
11
|
+
The model learns by comparing predictions with target tokens.
|
|
12
|
+
The loss measures how incorrect the predictions are.
|
|
13
|
+
Backpropagation calculates gradients.
|
|
14
|
+
The optimizer updates the model parameters.
|
|
15
|
+
Training helps the model improve its next token predictions.
|
|
16
|
+
A GPT model generates text one token at a time.
|
|
17
|
+
The generated token becomes part of the next input.
|
|
@@ -0,0 +1,135 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn.functional as F
|
|
3
|
+
|
|
4
|
+
from model.gpt import GPT
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
# -----------------------------
|
|
8
|
+
# Load checkpoint
|
|
9
|
+
# -----------------------------
|
|
10
|
+
|
|
11
|
+
checkpoint = torch.load(
|
|
12
|
+
"tokencraft_model.pth",
|
|
13
|
+
map_location="cpu",
|
|
14
|
+
weights_only=False
|
|
15
|
+
)
|
|
16
|
+
|
|
17
|
+
config = checkpoint["config"]
|
|
18
|
+
vocab = checkpoint["vocab"]
|
|
19
|
+
|
|
20
|
+
# Build reverse vocabulary
|
|
21
|
+
reverse_vocab = {
|
|
22
|
+
token_id: token
|
|
23
|
+
for token, token_id in vocab.items()
|
|
24
|
+
}
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
# -----------------------------
|
|
28
|
+
# Select device
|
|
29
|
+
# -----------------------------
|
|
30
|
+
|
|
31
|
+
if torch.backends.mps.is_available():
|
|
32
|
+
device = torch.device("mps")
|
|
33
|
+
else:
|
|
34
|
+
device = torch.device("cpu")
|
|
35
|
+
|
|
36
|
+
print("Using device:", device)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
# -----------------------------
|
|
40
|
+
# Rebuild model
|
|
41
|
+
# -----------------------------
|
|
42
|
+
|
|
43
|
+
model = GPT(
|
|
44
|
+
vocab_size=config["vocab_size"],
|
|
45
|
+
context_length=config["context_length"],
|
|
46
|
+
embedding_dim=config["embedding_dim"],
|
|
47
|
+
num_heads=config["num_heads"],
|
|
48
|
+
hidden_dim=config["hidden_dim"],
|
|
49
|
+
num_layers=config["num_layers"]
|
|
50
|
+
).to(device)
|
|
51
|
+
|
|
52
|
+
model.load_state_dict(checkpoint["model_state_dict"])
|
|
53
|
+
model.eval()
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
# -----------------------------
|
|
57
|
+
# Text generation
|
|
58
|
+
# -----------------------------
|
|
59
|
+
|
|
60
|
+
def generate_text(prompt, max_new_tokens=10):
|
|
61
|
+
# Convert prompt words into token IDs
|
|
62
|
+
tokens = prompt.split()
|
|
63
|
+
|
|
64
|
+
token_ids = [
|
|
65
|
+
vocab.get(token, vocab.get("<UNK>", 0))
|
|
66
|
+
for token in tokens
|
|
67
|
+
]
|
|
68
|
+
|
|
69
|
+
input_ids = torch.tensor(
|
|
70
|
+
[token_ids],
|
|
71
|
+
dtype=torch.long,
|
|
72
|
+
device=device
|
|
73
|
+
)
|
|
74
|
+
|
|
75
|
+
generated_ids = token_ids.copy()
|
|
76
|
+
|
|
77
|
+
with torch.no_grad():
|
|
78
|
+
for _ in range(max_new_tokens):
|
|
79
|
+
# Keep only the most recent context tokens
|
|
80
|
+
context = input_ids[:, -config["context_length"]:]
|
|
81
|
+
|
|
82
|
+
# Predict next token
|
|
83
|
+
logits = model(context)
|
|
84
|
+
|
|
85
|
+
# Get logits for the final position
|
|
86
|
+
next_token_logits = logits[:, -1, :]
|
|
87
|
+
|
|
88
|
+
# Apply temperature scaling
|
|
89
|
+
temperature = 0.8
|
|
90
|
+
scaled_logits = next_token_logits / temperature
|
|
91
|
+
|
|
92
|
+
# Convert scaled logits to probabilities
|
|
93
|
+
probabilities = F.softmax(
|
|
94
|
+
scaled_logits,
|
|
95
|
+
dim=-1
|
|
96
|
+
)
|
|
97
|
+
|
|
98
|
+
# Randomly sample the next token
|
|
99
|
+
next_token_id = torch.multinomial(
|
|
100
|
+
probabilities,
|
|
101
|
+
num_samples=1
|
|
102
|
+
).item()
|
|
103
|
+
|
|
104
|
+
generated_ids.append(next_token_id)
|
|
105
|
+
|
|
106
|
+
# Add prediction to input
|
|
107
|
+
input_ids = torch.cat(
|
|
108
|
+
[
|
|
109
|
+
input_ids,
|
|
110
|
+
torch.tensor(
|
|
111
|
+
[[next_token_id]],
|
|
112
|
+
dtype=torch.long,
|
|
113
|
+
device=device
|
|
114
|
+
)
|
|
115
|
+
],
|
|
116
|
+
dim=1
|
|
117
|
+
)
|
|
118
|
+
|
|
119
|
+
generated_tokens = [
|
|
120
|
+
reverse_vocab.get(token_id, "<UNK>")
|
|
121
|
+
for token_id in generated_ids
|
|
122
|
+
]
|
|
123
|
+
|
|
124
|
+
return " ".join(generated_tokens)
|
|
125
|
+
|
|
126
|
+
|
|
127
|
+
# -----------------------------
|
|
128
|
+
# Run generation
|
|
129
|
+
# -----------------------------
|
|
130
|
+
|
|
131
|
+
prompt = "the model"
|
|
132
|
+
result = generate_text(prompt, max_new_tokens=10)
|
|
133
|
+
|
|
134
|
+
print("\nPrompt:", prompt)
|
|
135
|
+
print("Generated text:", result)
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class TokenEmbedding(nn.Module):
|
|
6
|
+
def __init__(self, vocab_size, embedding_dim):
|
|
7
|
+
super().__init__()
|
|
8
|
+
|
|
9
|
+
self.embedding = nn.Embedding(
|
|
10
|
+
num_embeddings=vocab_size,
|
|
11
|
+
embedding_dim=embedding_dim
|
|
12
|
+
)
|
|
13
|
+
|
|
14
|
+
def forward(self, token_ids):
|
|
15
|
+
return self.embedding(token_ids)
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
if __name__ == "__main__":
|
|
19
|
+
vocab_size = 10
|
|
20
|
+
embedding_dim = 8
|
|
21
|
+
|
|
22
|
+
model = TokenEmbedding(vocab_size, embedding_dim)
|
|
23
|
+
|
|
24
|
+
token_ids = torch.tensor([4, 7, 2])
|
|
25
|
+
|
|
26
|
+
embeddings = model(token_ids)
|
|
27
|
+
|
|
28
|
+
print("Token IDs:")
|
|
29
|
+
print(token_ids)
|
|
30
|
+
|
|
31
|
+
print("\nEmbeddings:")
|
|
32
|
+
print(embeddings)
|
|
33
|
+
|
|
34
|
+
print("\nEmbedding shape:")
|
|
35
|
+
print(embeddings.shape)
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class FeedForward(nn.Module):
|
|
6
|
+
def __init__(self, embedding_dim, hidden_dim):
|
|
7
|
+
super().__init__()
|
|
8
|
+
|
|
9
|
+
self.network = nn.Sequential(
|
|
10
|
+
nn.Linear(embedding_dim, hidden_dim),
|
|
11
|
+
nn.GELU(),
|
|
12
|
+
nn.Linear(hidden_dim, embedding_dim)
|
|
13
|
+
)
|
|
14
|
+
|
|
15
|
+
def forward(self, x):
|
|
16
|
+
return self.network(x)
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
if __name__ == "__main__":
|
|
20
|
+
batch_size = 2
|
|
21
|
+
sequence_length = 5
|
|
22
|
+
embedding_dim = 16
|
|
23
|
+
hidden_dim = 64
|
|
24
|
+
|
|
25
|
+
x = torch.randn(
|
|
26
|
+
batch_size,
|
|
27
|
+
sequence_length,
|
|
28
|
+
embedding_dim
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
feed_forward = FeedForward(
|
|
32
|
+
embedding_dim=embedding_dim,
|
|
33
|
+
hidden_dim=hidden_dim
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
output = feed_forward(x)
|
|
37
|
+
|
|
38
|
+
print("Input shape:", x.shape)
|
|
39
|
+
print("Output shape:", output.shape)
|
|
@@ -0,0 +1,86 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
|
|
4
|
+
from .input_embedding import InputEmbedding
|
|
5
|
+
from .transformer_block import TransformerBlock
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class GPT(nn.Module):
|
|
9
|
+
def __init__(
|
|
10
|
+
self,
|
|
11
|
+
vocab_size,
|
|
12
|
+
context_length,
|
|
13
|
+
embedding_dim,
|
|
14
|
+
num_heads,
|
|
15
|
+
hidden_dim,
|
|
16
|
+
num_layers
|
|
17
|
+
):
|
|
18
|
+
super().__init__()
|
|
19
|
+
|
|
20
|
+
self.context_length = context_length
|
|
21
|
+
|
|
22
|
+
self.embedding = InputEmbedding(
|
|
23
|
+
vocab_size=vocab_size,
|
|
24
|
+
context_length=context_length,
|
|
25
|
+
embedding_dim=embedding_dim
|
|
26
|
+
)
|
|
27
|
+
|
|
28
|
+
self.transformer_blocks = nn.Sequential(
|
|
29
|
+
*[
|
|
30
|
+
TransformerBlock(
|
|
31
|
+
embedding_dim=embedding_dim,
|
|
32
|
+
num_heads=num_heads,
|
|
33
|
+
hidden_dim=hidden_dim
|
|
34
|
+
)
|
|
35
|
+
for _ in range(num_layers)
|
|
36
|
+
]
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
self.final_norm = nn.LayerNorm(embedding_dim)
|
|
40
|
+
|
|
41
|
+
self.output_head = nn.Linear(
|
|
42
|
+
embedding_dim,
|
|
43
|
+
vocab_size,
|
|
44
|
+
bias=False
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
def forward(self, token_ids):
|
|
48
|
+
x = self.embedding(token_ids)
|
|
49
|
+
|
|
50
|
+
x = self.transformer_blocks(x)
|
|
51
|
+
|
|
52
|
+
x = self.final_norm(x)
|
|
53
|
+
|
|
54
|
+
logits = self.output_head(x)
|
|
55
|
+
|
|
56
|
+
return logits
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
if __name__ == "__main__":
|
|
60
|
+
vocab_size = 20
|
|
61
|
+
context_length = 5
|
|
62
|
+
embedding_dim = 16
|
|
63
|
+
num_heads = 4
|
|
64
|
+
hidden_dim = 64
|
|
65
|
+
num_layers = 2
|
|
66
|
+
|
|
67
|
+
token_ids = torch.randint(
|
|
68
|
+
0,
|
|
69
|
+
vocab_size,
|
|
70
|
+
(2, context_length)
|
|
71
|
+
)
|
|
72
|
+
|
|
73
|
+
model = GPT(
|
|
74
|
+
vocab_size=vocab_size,
|
|
75
|
+
context_length=context_length,
|
|
76
|
+
embedding_dim=embedding_dim,
|
|
77
|
+
num_heads=num_heads,
|
|
78
|
+
hidden_dim=hidden_dim,
|
|
79
|
+
num_layers=num_layers
|
|
80
|
+
)
|
|
81
|
+
|
|
82
|
+
logits = model(token_ids)
|
|
83
|
+
|
|
84
|
+
print("Input shape:", token_ids.shape)
|
|
85
|
+
print("Logits shape:", logits.shape)
|
|
86
|
+
print("Parameter count:", sum(p.numel() for p in model.parameters()))
|
|
@@ -0,0 +1,54 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class InputEmbedding(nn.Module):
|
|
6
|
+
def __init__(self, vocab_size, context_length, embedding_dim):
|
|
7
|
+
super().__init__()
|
|
8
|
+
|
|
9
|
+
self.token_embedding = nn.Embedding(
|
|
10
|
+
num_embeddings=vocab_size,
|
|
11
|
+
embedding_dim=embedding_dim
|
|
12
|
+
)
|
|
13
|
+
|
|
14
|
+
self.position_embedding = nn.Embedding(
|
|
15
|
+
num_embeddings=context_length,
|
|
16
|
+
embedding_dim=embedding_dim
|
|
17
|
+
)
|
|
18
|
+
|
|
19
|
+
def forward(self, token_ids):
|
|
20
|
+
sequence_length = token_ids.shape[1]
|
|
21
|
+
|
|
22
|
+
positions = torch.arange(
|
|
23
|
+
sequence_length,
|
|
24
|
+
device=token_ids.device
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
token_vectors = self.token_embedding(token_ids)
|
|
28
|
+
position_vectors = self.position_embedding(positions)
|
|
29
|
+
|
|
30
|
+
return token_vectors + position_vectors
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
if __name__ == "__main__":
|
|
34
|
+
vocab_size = 20
|
|
35
|
+
context_length = 5
|
|
36
|
+
embedding_dim = 8
|
|
37
|
+
|
|
38
|
+
token_ids = torch.tensor([
|
|
39
|
+
[1, 2, 3, 4, 5],
|
|
40
|
+
[5, 4, 3, 2, 1]
|
|
41
|
+
])
|
|
42
|
+
|
|
43
|
+
model = InputEmbedding(
|
|
44
|
+
vocab_size=vocab_size,
|
|
45
|
+
context_length=context_length,
|
|
46
|
+
embedding_dim=embedding_dim
|
|
47
|
+
)
|
|
48
|
+
|
|
49
|
+
output = model(token_ids)
|
|
50
|
+
|
|
51
|
+
print("Output:")
|
|
52
|
+
print(output)
|
|
53
|
+
|
|
54
|
+
print("Shape:", output.shape)
|
|
@@ -0,0 +1,117 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
import math
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
class MultiHeadAttention(nn.Module):
|
|
7
|
+
def __init__(self, embedding_dim, num_heads):
|
|
8
|
+
super().__init__()
|
|
9
|
+
|
|
10
|
+
if embedding_dim % num_heads != 0:
|
|
11
|
+
raise ValueError(
|
|
12
|
+
"embedding_dim must be divisible by num_heads"
|
|
13
|
+
)
|
|
14
|
+
|
|
15
|
+
self.embedding_dim = embedding_dim
|
|
16
|
+
self.num_heads = num_heads
|
|
17
|
+
self.head_dim = embedding_dim // num_heads
|
|
18
|
+
|
|
19
|
+
self.query = nn.Linear(embedding_dim, embedding_dim)
|
|
20
|
+
self.key = nn.Linear(embedding_dim, embedding_dim)
|
|
21
|
+
self.value = nn.Linear(embedding_dim, embedding_dim)
|
|
22
|
+
|
|
23
|
+
self.output_projection = nn.Linear(
|
|
24
|
+
embedding_dim,
|
|
25
|
+
embedding_dim
|
|
26
|
+
)
|
|
27
|
+
|
|
28
|
+
self.scale = math.sqrt(self.head_dim)
|
|
29
|
+
|
|
30
|
+
def forward(self, x):
|
|
31
|
+
batch_size, sequence_length, _ = x.shape
|
|
32
|
+
|
|
33
|
+
Q = self.query(x)
|
|
34
|
+
K = self.key(x)
|
|
35
|
+
V = self.value(x)
|
|
36
|
+
|
|
37
|
+
# Split into multiple heads
|
|
38
|
+
Q = Q.view(
|
|
39
|
+
batch_size,
|
|
40
|
+
sequence_length,
|
|
41
|
+
self.num_heads,
|
|
42
|
+
self.head_dim
|
|
43
|
+
).transpose(1, 2)
|
|
44
|
+
|
|
45
|
+
K = K.view(
|
|
46
|
+
batch_size,
|
|
47
|
+
sequence_length,
|
|
48
|
+
self.num_heads,
|
|
49
|
+
self.head_dim
|
|
50
|
+
).transpose(1, 2)
|
|
51
|
+
|
|
52
|
+
V = V.view(
|
|
53
|
+
batch_size,
|
|
54
|
+
sequence_length,
|
|
55
|
+
self.num_heads,
|
|
56
|
+
self.head_dim
|
|
57
|
+
).transpose(1, 2)
|
|
58
|
+
|
|
59
|
+
# Attention scores
|
|
60
|
+
scores = torch.matmul(
|
|
61
|
+
Q,
|
|
62
|
+
K.transpose(-2, -1)
|
|
63
|
+
) / self.scale
|
|
64
|
+
|
|
65
|
+
# Causal mask
|
|
66
|
+
mask = torch.triu(
|
|
67
|
+
torch.ones(
|
|
68
|
+
sequence_length,
|
|
69
|
+
sequence_length,
|
|
70
|
+
device=x.device,
|
|
71
|
+
dtype=torch.bool
|
|
72
|
+
),
|
|
73
|
+
diagonal=1
|
|
74
|
+
)
|
|
75
|
+
|
|
76
|
+
scores = scores.masked_fill(mask, float("-inf"))
|
|
77
|
+
|
|
78
|
+
attention_weights = torch.softmax(
|
|
79
|
+
scores,
|
|
80
|
+
dim=-1
|
|
81
|
+
)
|
|
82
|
+
|
|
83
|
+
output = torch.matmul(attention_weights, V)
|
|
84
|
+
|
|
85
|
+
# Combine attention heads
|
|
86
|
+
output = output.transpose(1, 2).contiguous()
|
|
87
|
+
|
|
88
|
+
output = output.view(
|
|
89
|
+
batch_size,
|
|
90
|
+
sequence_length,
|
|
91
|
+
self.embedding_dim
|
|
92
|
+
)
|
|
93
|
+
|
|
94
|
+
return self.output_projection(output)
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
if __name__ == "__main__":
|
|
98
|
+
batch_size = 2
|
|
99
|
+
sequence_length = 5
|
|
100
|
+
embedding_dim = 16
|
|
101
|
+
num_heads = 4
|
|
102
|
+
|
|
103
|
+
x = torch.randn(
|
|
104
|
+
batch_size,
|
|
105
|
+
sequence_length,
|
|
106
|
+
embedding_dim
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
attention = MultiHeadAttention(
|
|
110
|
+
embedding_dim=embedding_dim,
|
|
111
|
+
num_heads=num_heads
|
|
112
|
+
)
|
|
113
|
+
|
|
114
|
+
output = attention(x)
|
|
115
|
+
|
|
116
|
+
print("Input shape:", x.shape)
|
|
117
|
+
print("Output shape:", output.shape)
|
|
@@ -0,0 +1,45 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class PositionalEmbedding(nn.Module):
|
|
6
|
+
def __init__(self, context_length, embedding_dim):
|
|
7
|
+
super().__init__()
|
|
8
|
+
|
|
9
|
+
self.position_embedding = nn.Embedding(
|
|
10
|
+
num_embeddings=context_length,
|
|
11
|
+
embedding_dim=embedding_dim
|
|
12
|
+
)
|
|
13
|
+
|
|
14
|
+
def forward(self, token_ids):
|
|
15
|
+
sequence_length = token_ids.shape[1]
|
|
16
|
+
|
|
17
|
+
positions = torch.arange(
|
|
18
|
+
sequence_length,
|
|
19
|
+
device=token_ids.device
|
|
20
|
+
)
|
|
21
|
+
|
|
22
|
+
return self.position_embedding(positions)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
if __name__ == "__main__":
|
|
26
|
+
batch_size = 2
|
|
27
|
+
context_length = 5
|
|
28
|
+
embedding_dim = 8
|
|
29
|
+
|
|
30
|
+
token_ids = torch.tensor([
|
|
31
|
+
[1, 2, 3, 4, 5],
|
|
32
|
+
[5, 4, 3, 2, 1]
|
|
33
|
+
])
|
|
34
|
+
|
|
35
|
+
model = PositionalEmbedding(
|
|
36
|
+
context_length=context_length,
|
|
37
|
+
embedding_dim=embedding_dim
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
positions = model(token_ids)
|
|
41
|
+
|
|
42
|
+
print("Positional embeddings:")
|
|
43
|
+
print(positions)
|
|
44
|
+
|
|
45
|
+
print("Shape:", positions.shape)
|
|
@@ -0,0 +1,64 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
import math
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
class SelfAttention(nn.Module):
|
|
7
|
+
def __init__(self, embedding_dim):
|
|
8
|
+
super().__init__()
|
|
9
|
+
|
|
10
|
+
self.query = nn.Linear(embedding_dim, embedding_dim)
|
|
11
|
+
self.key = nn.Linear(embedding_dim, embedding_dim)
|
|
12
|
+
self.value = nn.Linear(embedding_dim, embedding_dim)
|
|
13
|
+
|
|
14
|
+
self.scale = math.sqrt(embedding_dim)
|
|
15
|
+
|
|
16
|
+
def forward(self, x):
|
|
17
|
+
# x shape: (batch_size, sequence_length, embedding_dim)
|
|
18
|
+
|
|
19
|
+
Q = self.query(x)
|
|
20
|
+
K = self.key(x)
|
|
21
|
+
V = self.value(x)
|
|
22
|
+
|
|
23
|
+
# Calculate attention scores
|
|
24
|
+
scores = torch.matmul(Q, K.transpose(-2, -1))
|
|
25
|
+
scores = scores / self.scale
|
|
26
|
+
|
|
27
|
+
# Causal mask: prevent looking at future tokens
|
|
28
|
+
sequence_length = x.shape[1]
|
|
29
|
+
|
|
30
|
+
mask = torch.triu(
|
|
31
|
+
torch.ones(
|
|
32
|
+
sequence_length,
|
|
33
|
+
sequence_length,
|
|
34
|
+
device=x.device,
|
|
35
|
+
dtype=torch.bool
|
|
36
|
+
),
|
|
37
|
+
diagonal=1
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
scores = scores.masked_fill(mask, float("-inf"))
|
|
41
|
+
|
|
42
|
+
# Convert scores into probabilities
|
|
43
|
+
attention_weights = torch.softmax(scores, dim=-1)
|
|
44
|
+
|
|
45
|
+
# Combine information from previous tokens
|
|
46
|
+
output = torch.matmul(attention_weights, V)
|
|
47
|
+
|
|
48
|
+
return output
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
if __name__ == "__main__":
|
|
52
|
+
batch_size = 2
|
|
53
|
+
sequence_length = 5
|
|
54
|
+
embedding_dim = 8
|
|
55
|
+
|
|
56
|
+
x = torch.randn(batch_size, sequence_length, embedding_dim)
|
|
57
|
+
|
|
58
|
+
attention = SelfAttention(embedding_dim)
|
|
59
|
+
|
|
60
|
+
output = attention(x)
|
|
61
|
+
|
|
62
|
+
print("Output shape:", output.shape)
|
|
63
|
+
print("Output:")
|
|
64
|
+
print(output)
|