-
Notifications
You must be signed in to change notification settings - Fork 0
/
bigram.py
130 lines (98 loc) · 3.67 KB
/
bigram.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import torch
import torch.nn as nn
import torch.nn.functional as F
batch_size = 32
block_size = 8
max_iters = 3000
eval_interval = 300
learning_rate = 1e-2
torch.manual_seed(1337)
eval_iters = 200
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("Running on device: ", device)
# !wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
#read to inspect
with open('input.txt','r',encoding='utf-8') as f:
text = f.read()
#unique characters in the dataset
chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = {ch:i for i,ch in enumerate(chars)}
itos ={i:ch for i,ch in enumerate(chars)}
encode = lambda s: [stoi[ch] for ch in s] ##encoder
decode = lambda l: ''.join([itos[i] for i in l]) ##decoder
#encoding entire dataset
data = torch.tensor(encode(text),dtype=torch.long)
print(data.shape,data.dtype)
#train val split
n = int(0.9*len(data))
train_data, val_data = data[:n], data[n:]
def get_batch(split):
data = train_data if split == 'train' else val_data
ix = torch.randint(len(data)-block_size,(batch_size,))
x = torch.stack([data[i:i+block_size] for i in ix])
y= torch.stack([data[i+1:i+block_size+1] for i in ix])
x,y = x.to(device),y.to(device)
return x,y
#cntext manager
@torch.no_grad()
def estimate_loss():
out={}
model.eval()
for split in ['train','val']:
losses = torch.zeros(eval_iters)
for k in range(eval_iters):
X,Y = get_batch(split)
logits,loss = model(X,Y)
losses[k] = loss.item()
out[split] = losses.mean()
model.train()
return out
class BigramLanguageModel(nn.Module):
def __init__(self,vocab_size):
super().__init__()
#each token direclty reads off the logits for the next token from a lookup table
self.token_embedding_table = nn.Embedding(vocab_size,vocab_size)
def forward(self,idx,targets =None):
#idx and targets are both [block_size,batch_size]
logits = self.token_embedding_table(idx) # [block_size,batch_size,vocab_size] (B,T,C)
if targets is None:
loss=None
else:
B,T,C = logits.shape
logits = logits.view(B*T,C)
targets = targets.view(B*T)
loss = F.cross_entropy(logits,targets)
return logits, loss
def generate(self, idx, max_new_tokens):
# idx is (B, T) array of indices in the current context
for _ in range(max_new_tokens):
# get the predictions
logits, loss = self(idx)
# focus only on the last time step
logits = logits[:, -1, :] # becomes (B, C)
# apply softmax to get probabilities
probs = F.softmax(logits, dim=-1) # (B, C)
# sample from the distribution
idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
# append sampled index to the running sequence
idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
return idx
model = BigramLanguageModel(vocab_size)
m= model.to(device)
optimizer = torch.optim.AdamW(m.parameters(),lr = learning_rate)
#create an optimizer
for iter in range(max_iters):
#calculating error over eval_iters
if iter % eval_iters == 0:
losses = estimate_loss()
print(f'Step {iter:5d}, Train loss: {losses["train"]:.4f}, Val loss: {losses["val"]:.4f}')
#sample a batch of data
xb,yb = get_batch('train')
#evalute the loss
logits,loss = m(xb,yb)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
context = torch.zeros((1,1),dtype=torch.long,device=device)
print(decode( m.generate(context,max_new_tokens=500)[0].tolist()))