piyushgrover commited on
Commit
d39406e
·
1 Parent(s): 7f8f59f

added files

Browse files
Files changed (5) hide show
  1. app.py +77 -0
  2. config.py +35 -0
  3. model.py +330 -0
  4. requirements.txt +7 -0
  5. test.ipynb +254 -0
app.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import gradio as gr
3
+ import random
4
+ from config import device_type, ckpt_path, GPTConfig, GPT, encode, decode, ctx, num_samples, max_new_tokens, temperature, top_k
5
+
6
+ checkpoint = torch.load(ckpt_path, map_location=device_type)
7
+ gptconf = GPTConfig(**checkpoint['model_args'])
8
+ model = GPT(gptconf)
9
+ state_dict = checkpoint['model']
10
+ unwanted_prefix = '_orig_mod.'
11
+ for k,v in list(state_dict.items()):
12
+ if k.startswith(unwanted_prefix):
13
+ state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
14
+ model.load_state_dict(state_dict)
15
+ model.eval()
16
+ model.to(device_type)
17
+
18
+
19
+ def fn_query_on_load():
20
+ return "in the air and"
21
+
22
+
23
+ def generate_commentary(start):
24
+ start_ids = encode(start)
25
+ x = (torch.tensor(start_ids, dtype=torch.long, device=device_type)[None, ...])
26
+
27
+ out_text = ''
28
+ with torch.no_grad():
29
+ with ctx:
30
+ for k in range(num_samples):
31
+ y = model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)
32
+ out_text += decode(y[0].tolist())
33
+ out_text += '\n-o-o-o-o-o-o-o-\n\n'
34
+
35
+ return {
36
+ output: out_text
37
+ }
38
+
39
+
40
+ with gr.Blocks() as app:
41
+ with gr.Row():
42
+ gr.Markdown(
43
+ """
44
+ # NanoGPT - Cricket Commentary Generative AI
45
+ ### Give a prompt and see how it comes out with cricket commentary :)
46
+ """)
47
+
48
+ with gr.Row(visible=True):
49
+ search_text = gr.Textbox(value=fn_query_on_load, placeholder='Enter prompt..', label='Enter Prompt')
50
+
51
+ with gr.Row():
52
+ submit_btn = gr.Button("Submit", variant='primary')
53
+ clear_btn = gr.ClearButton()
54
+ with gr.Row():
55
+ with gr.Row():
56
+ output = gr.Textbox(lines=15, interactive=False, label='Commentary Box')
57
+
58
+ def clear_data():
59
+ return {
60
+ output: None,
61
+ search_text: None
62
+ }
63
+
64
+ clear_btn.click(clear_data, None, [output, search_text])
65
+
66
+
67
+ submit_btn.click(
68
+ generate_commentary,
69
+ search_text,
70
+ output
71
+ )
72
+
73
+
74
+ '''
75
+ Launch the app
76
+ '''
77
+ app.queue().launch()
config.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import pickle
3
+ from contextlib import nullcontext
4
+ import torch
5
+ import tiktoken
6
+ from model import GPTConfig, GPT
7
+ import random
8
+
9
+ out_dir = 'out-cricket-commentary' # ignored if init_from is not 'resume'
10
+ start = "\n" # or "<|endoftext|>" or etc. Can also specify a file, use as: "FILE:prompt.txt"
11
+ num_samples = 1 #10 # number of samples to draw
12
+ max_new_tokens = 300 #500 # number of tokens generated in each sample
13
+ temperature = 0.8 # 1.0 = no change, < 1.0 = less random, > 1.0 = more random, in predictions
14
+ top_k = 50 #200 # retain only the top_k most likely tokens, clamp others to have 0 probability
15
+ seed = random.randint(1, 10000)
16
+ device = 'cuda' if torch.cuda.is_available() else 'cpu' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc.
17
+ dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 'float32' or 'bfloat16' or 'float16'
18
+ compile = False # use PyTorch 2.0 to compile the model to be faster
19
+ #exec(open('configurator.py').read()) # overrides from command line or config file
20
+ # -----------------------------------------------------------------------------
21
+
22
+ torch.manual_seed(seed)
23
+ torch.cuda.manual_seed(seed)
24
+ torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
25
+ torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
26
+ device_type = 'cuda' if 'cuda' in device else 'cpu' # for later use in torch.autocast
27
+ ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
28
+ ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
29
+
30
+ ckpt_path = os.path.join(out_dir, 'ckpt.pt')
31
+
32
+ enc = tiktoken.get_encoding("gpt2")
33
+ encode = lambda s: enc.encode(s, allowed_special={"<|endoftext|>"})
34
+ decode = lambda l: enc.decode(l)
35
+
model.py ADDED
@@ -0,0 +1,330 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Full definition of a GPT Language Model, all of it in this single file.
3
+ References:
4
+ 1) the official GPT-2 TensorFlow implementation released by OpenAI:
5
+ https://github.com/openai/gpt-2/blob/master/src/model.py
6
+ 2) huggingface/transformers PyTorch implementation:
7
+ https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
8
+ """
9
+
10
+ import math
11
+ import inspect
12
+ from dataclasses import dataclass
13
+
14
+ import torch
15
+ import torch.nn as nn
16
+ from torch.nn import functional as F
17
+
18
+ class LayerNorm(nn.Module):
19
+ """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """
20
+
21
+ def __init__(self, ndim, bias):
22
+ super().__init__()
23
+ self.weight = nn.Parameter(torch.ones(ndim))
24
+ self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
25
+
26
+ def forward(self, input):
27
+ return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
28
+
29
+ class CausalSelfAttention(nn.Module):
30
+
31
+ def __init__(self, config):
32
+ super().__init__()
33
+ assert config.n_embd % config.n_head == 0
34
+ # key, query, value projections for all heads, but in a batch
35
+ self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
36
+ # output projection
37
+ self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
38
+ # regularization
39
+ self.attn_dropout = nn.Dropout(config.dropout)
40
+ self.resid_dropout = nn.Dropout(config.dropout)
41
+ self.n_head = config.n_head
42
+ self.n_embd = config.n_embd
43
+ self.dropout = config.dropout
44
+ # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
45
+ self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
46
+ if not self.flash:
47
+ print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
48
+ # causal mask to ensure that attention is only applied to the left in the input sequence
49
+ self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
50
+ .view(1, 1, config.block_size, config.block_size))
51
+
52
+ def forward(self, x):
53
+ B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
54
+
55
+ # calculate query, key, values for all heads in batch and move head forward to be the batch dim
56
+ q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
57
+ k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
58
+ q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
59
+ v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
60
+
61
+ # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
62
+ if self.flash:
63
+ # efficient attention using Flash Attention CUDA kernels
64
+ y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
65
+ else:
66
+ # manual implementation of attention
67
+ att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
68
+ att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
69
+ att = F.softmax(att, dim=-1)
70
+ att = self.attn_dropout(att)
71
+ y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
72
+ y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
73
+
74
+ # output projection
75
+ y = self.resid_dropout(self.c_proj(y))
76
+ return y
77
+
78
+ class MLP(nn.Module):
79
+
80
+ def __init__(self, config):
81
+ super().__init__()
82
+ self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
83
+ self.gelu = nn.GELU()
84
+ self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
85
+ self.dropout = nn.Dropout(config.dropout)
86
+
87
+ def forward(self, x):
88
+ x = self.c_fc(x)
89
+ x = self.gelu(x)
90
+ x = self.c_proj(x)
91
+ x = self.dropout(x)
92
+ return x
93
+
94
+ class Block(nn.Module):
95
+
96
+ def __init__(self, config):
97
+ super().__init__()
98
+ self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
99
+ self.attn = CausalSelfAttention(config)
100
+ self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
101
+ self.mlp = MLP(config)
102
+
103
+ def forward(self, x):
104
+ x = x + self.attn(self.ln_1(x))
105
+ x = x + self.mlp(self.ln_2(x))
106
+ return x
107
+
108
+ @dataclass
109
+ class GPTConfig:
110
+ block_size: int = 1024
111
+ vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
112
+ n_layer: int = 12
113
+ n_head: int = 12
114
+ n_embd: int = 768
115
+ dropout: float = 0.0
116
+ bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
117
+
118
+ class GPT(nn.Module):
119
+
120
+ def __init__(self, config):
121
+ super().__init__()
122
+ assert config.vocab_size is not None
123
+ assert config.block_size is not None
124
+ self.config = config
125
+
126
+ self.transformer = nn.ModuleDict(dict(
127
+ wte = nn.Embedding(config.vocab_size, config.n_embd),
128
+ wpe = nn.Embedding(config.block_size, config.n_embd),
129
+ drop = nn.Dropout(config.dropout),
130
+ h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
131
+ ln_f = LayerNorm(config.n_embd, bias=config.bias),
132
+ ))
133
+ self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
134
+ # with weight tying when using torch.compile() some warnings get generated:
135
+ # "UserWarning: functional_call was passed multiple values for tied weights.
136
+ # This behavior is deprecated and will be an error in future versions"
137
+ # not 100% sure what this is, so far seems to be harmless. TODO investigate
138
+ self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
139
+
140
+ # init all weights
141
+ self.apply(self._init_weights)
142
+ # apply special scaled init to the residual projections, per GPT-2 paper
143
+ for pn, p in self.named_parameters():
144
+ if pn.endswith('c_proj.weight'):
145
+ torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
146
+
147
+ # report number of parameters
148
+ print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))
149
+
150
+ def get_num_params(self, non_embedding=True):
151
+ """
152
+ Return the number of parameters in the model.
153
+ For non-embedding count (default), the position embeddings get subtracted.
154
+ The token embeddings would too, except due to the parameter sharing these
155
+ params are actually used as weights in the final layer, so we include them.
156
+ """
157
+ n_params = sum(p.numel() for p in self.parameters())
158
+ if non_embedding:
159
+ n_params -= self.transformer.wpe.weight.numel()
160
+ return n_params
161
+
162
+ def _init_weights(self, module):
163
+ if isinstance(module, nn.Linear):
164
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
165
+ if module.bias is not None:
166
+ torch.nn.init.zeros_(module.bias)
167
+ elif isinstance(module, nn.Embedding):
168
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
169
+
170
+ def forward(self, idx, targets=None):
171
+ device = idx.device
172
+ b, t = idx.size()
173
+ assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
174
+ pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)
175
+
176
+ # forward the GPT model itself
177
+ tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
178
+ pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
179
+ x = self.transformer.drop(tok_emb + pos_emb)
180
+ for block in self.transformer.h:
181
+ x = block(x)
182
+ x = self.transformer.ln_f(x)
183
+
184
+ if targets is not None:
185
+ # if we are given some desired targets also calculate the loss
186
+ logits = self.lm_head(x)
187
+ loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
188
+ else:
189
+ # inference-time mini-optimization: only forward the lm_head on the very last position
190
+ logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
191
+ loss = None
192
+
193
+ return logits, loss
194
+
195
+ def crop_block_size(self, block_size):
196
+ # model surgery to decrease the block size if necessary
197
+ # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
198
+ # but want to use a smaller block size for some smaller, simpler model
199
+ assert block_size <= self.config.block_size
200
+ self.config.block_size = block_size
201
+ self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
202
+ for block in self.transformer.h:
203
+ if hasattr(block.attn, 'bias'):
204
+ block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]
205
+
206
+ @classmethod
207
+ def from_pretrained(cls, model_type, override_args=None):
208
+ assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
209
+ override_args = override_args or {} # default to empty dict
210
+ # only dropout can be overridden see more notes below
211
+ assert all(k == 'dropout' for k in override_args)
212
+ from transformers import GPT2LMHeadModel
213
+ print("loading weights from pretrained gpt: %s" % model_type)
214
+
215
+ # n_layer, n_head and n_embd are determined from model_type
216
+ config_args = {
217
+ 'gpt2': dict(n_layer=12, n_head=12, n_embd=768), # 124M params
218
+ 'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
219
+ 'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
220
+ 'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
221
+ }[model_type]
222
+ print("forcing vocab_size=50257, block_size=1024, bias=True")
223
+ config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
224
+ config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
225
+ config_args['bias'] = True # always True for GPT model checkpoints
226
+ # we can override the dropout rate, if desired
227
+ if 'dropout' in override_args:
228
+ print(f"overriding dropout rate to {override_args['dropout']}")
229
+ config_args['dropout'] = override_args['dropout']
230
+ # create a from-scratch initialized minGPT model
231
+ config = GPTConfig(**config_args)
232
+ model = GPT(config)
233
+ sd = model.state_dict()
234
+ sd_keys = sd.keys()
235
+ sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param
236
+
237
+ # init a huggingface/transformers model
238
+ model_hf = GPT2LMHeadModel.from_pretrained(model_type)
239
+ sd_hf = model_hf.state_dict()
240
+
241
+ # copy while ensuring all of the parameters are aligned and match in names and shapes
242
+ sd_keys_hf = sd_hf.keys()
243
+ sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer
244
+ sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)
245
+ transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
246
+ # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear
247
+ # this means that we have to transpose these weights when we import them
248
+ assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
249
+ for k in sd_keys_hf:
250
+ if any(k.endswith(w) for w in transposed):
251
+ # special treatment for the Conv1D weights we need to transpose
252
+ assert sd_hf[k].shape[::-1] == sd[k].shape
253
+ with torch.no_grad():
254
+ sd[k].copy_(sd_hf[k].t())
255
+ else:
256
+ # vanilla copy over the other parameters
257
+ assert sd_hf[k].shape == sd[k].shape
258
+ with torch.no_grad():
259
+ sd[k].copy_(sd_hf[k])
260
+
261
+ return model
262
+
263
+ def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
264
+ # start with all of the candidate parameters
265
+ param_dict = {pn: p for pn, p in self.named_parameters()}
266
+ # filter out those that do not require grad
267
+ param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
268
+ # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
269
+ # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
270
+ decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
271
+ nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
272
+ optim_groups = [
273
+ {'params': decay_params, 'weight_decay': weight_decay},
274
+ {'params': nodecay_params, 'weight_decay': 0.0}
275
+ ]
276
+ num_decay_params = sum(p.numel() for p in decay_params)
277
+ num_nodecay_params = sum(p.numel() for p in nodecay_params)
278
+ print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
279
+ print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
280
+ # Create AdamW optimizer and use the fused version if it is available
281
+ fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
282
+ use_fused = fused_available and device_type == 'cuda'
283
+ extra_args = dict(fused=True) if use_fused else dict()
284
+ optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
285
+ print(f"using fused AdamW: {use_fused}")
286
+
287
+ return optimizer
288
+
289
+ def estimate_mfu(self, fwdbwd_per_iter, dt):
290
+ """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
291
+ # first estimate the number of flops we do per iteration.
292
+ # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
293
+ N = self.get_num_params()
294
+ cfg = self.config
295
+ L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
296
+ flops_per_token = 6*N + 12*L*H*Q*T
297
+ flops_per_fwdbwd = flops_per_token * T
298
+ flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
299
+ # express our flops throughput as ratio of A100 bfloat16 peak flops
300
+ flops_achieved = flops_per_iter * (1.0/dt) # per second
301
+ flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
302
+ mfu = flops_achieved / flops_promised
303
+ return mfu
304
+
305
+ @torch.no_grad()
306
+ def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
307
+ """
308
+ Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
309
+ the sequence max_new_tokens times, feeding the predictions back into the model each time.
310
+ Most likely you'll want to make sure to be in model.eval() mode of operation for this.
311
+ """
312
+ for _ in range(max_new_tokens):
313
+ # if the sequence context is growing too long we must crop it at block_size
314
+ idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
315
+ # forward the model to get the logits for the index in the sequence
316
+ logits, _ = self(idx_cond)
317
+ # pluck the logits at the final step and scale by desired temperature
318
+ logits = logits[:, -1, :] / temperature
319
+ # optionally crop the logits to only the top k options
320
+ if top_k is not None:
321
+ v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
322
+ logits[logits < v[:, [-1]]] = -float('Inf')
323
+ # apply softmax to convert logits to (normalized) probabilities
324
+ probs = F.softmax(logits, dim=-1)
325
+ # sample from the distribution
326
+ idx_next = torch.multinomial(probs, num_samples=1)
327
+ # append sampled index to the running sequence and continue
328
+ idx = torch.cat((idx, idx_next), dim=1)
329
+
330
+ return idx
requirements.txt ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ torch
2
+ numpy
3
+ transformers
4
+ datasets
5
+ tiktoken
6
+ wandb
7
+ tqdm
test.ipynb ADDED
@@ -0,0 +1,254 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "cells": [
3
+ {
4
+ "cell_type": "code",
5
+ "execution_count": 1,
6
+ "metadata": {},
7
+ "outputs": [
8
+ {
9
+ "name": "stdout",
10
+ "output_type": "stream",
11
+ "text": [
12
+ "number of parameters: 29.94M\n",
13
+ "Running on local URL: http://127.0.0.1:7860\n",
14
+ "\n",
15
+ "To create a public link, set `share=True` in `launch()`.\n"
16
+ ]
17
+ },
18
+ {
19
+ "data": {
20
+ "text/html": [
21
+ "<div><iframe src=\"http://127.0.0.1:7860/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
22
+ ],
23
+ "text/plain": [
24
+ "<IPython.core.display.HTML object>"
25
+ ]
26
+ },
27
+ "metadata": {},
28
+ "output_type": "execute_result"
29
+ },
30
+ {
31
+ "data": {
32
+ "text/plain": []
33
+ },
34
+ "execution_count": 1,
35
+ "metadata": {},
36
+ "output_type": "execute_result"
37
+ }
38
+ ],
39
+ "source": [
40
+ "import torch\n",
41
+ "import gradio as gr\n",
42
+ "import random\n",
43
+ "from config import device_type, ckpt_path, GPTConfig, GPT, encode, decode, ctx, num_samples, max_new_tokens, temperature, top_k\n",
44
+ "\n",
45
+ "checkpoint = torch.load(ckpt_path, map_location=device_type)\n",
46
+ "gptconf = GPTConfig(**checkpoint['model_args'])\n",
47
+ "model = GPT(gptconf)\n",
48
+ "state_dict = checkpoint['model']\n",
49
+ "unwanted_prefix = '_orig_mod.'\n",
50
+ "for k,v in list(state_dict.items()):\n",
51
+ " if k.startswith(unwanted_prefix):\n",
52
+ " state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)\n",
53
+ "model.load_state_dict(state_dict)\n",
54
+ "model.eval()\n",
55
+ "model.to(device_type)\n",
56
+ "\n",
57
+ "button_click = False\n",
58
+ "\n",
59
+ "def fn_query_on_load():\n",
60
+ " return \"in the air and\"\n",
61
+ "\n",
62
+ "num_samples = 1\n",
63
+ "def generate_commentary(start):\n",
64
+ " start_ids = encode(start)\n",
65
+ " x = (torch.tensor(start_ids, dtype=torch.long, device=device_type)[None, ...])\n",
66
+ "\n",
67
+ " out_text = ''\n",
68
+ " with torch.no_grad():\n",
69
+ " with ctx:\n",
70
+ " for k in range(num_samples):\n",
71
+ " y = model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)\n",
72
+ " out_text += decode(y[0].tolist())\n",
73
+ " out_text += '\\n-o-o-o-o-o-o-o-\\n\\n'\n",
74
+ "\n",
75
+ " return out_text\n",
76
+ " \n",
77
+ " \n",
78
+ "def fn_gen_comm(prompt, st, o1, o2, o3):\n",
79
+ " '''global button_click\n",
80
+ " if not button_click:\n",
81
+ " button_click = True\n",
82
+ " elif stat == -1:\n",
83
+ " button_click = False\n",
84
+ " return {\n",
85
+ " output1: output1,\n",
86
+ " output2: output2,\n",
87
+ " output3: output3,\n",
88
+ " stat: stat\n",
89
+ " }\n",
90
+ " \n",
91
+ " \n",
92
+ " out = generate_commentary(prompt)\n",
93
+ " if stat == -1:\n",
94
+ " return {\n",
95
+ " output1: out,\n",
96
+ " output2: None,\n",
97
+ " output3: None,\n",
98
+ " stat: 0\n",
99
+ " }\n",
100
+ " \n",
101
+ " elif stat == 0:\n",
102
+ " return {\n",
103
+ " output1: output1,\n",
104
+ " output2: out,\n",
105
+ " output3: None,\n",
106
+ " stat: 1\n",
107
+ " }\n",
108
+ " \n",
109
+ " elif stat == 2:\n",
110
+ " return {\n",
111
+ " output1: output1,\n",
112
+ " output2: output2,\n",
113
+ " output3: out,\n",
114
+ " stat: -1\n",
115
+ " }'''\n",
116
+ " \n",
117
+ " global button_click\n",
118
+ " if not button_click:\n",
119
+ " if st == -1:\n",
120
+ " button_click = True\n",
121
+ " elif st == -1:\n",
122
+ " button_click = False\n",
123
+ " return {\n",
124
+ " output1: o1,\n",
125
+ " output2: o2,\n",
126
+ " output3: o3,\n",
127
+ " stat: -1\n",
128
+ " }\n",
129
+ " elif st == 2:\n",
130
+ " button_click = False\n",
131
+ " return {\n",
132
+ " output1: o1,\n",
133
+ " output2: o2,\n",
134
+ " output3: o3,\n",
135
+ " stat: -1\n",
136
+ " }\n",
137
+ " \n",
138
+ " out = generate_commentary(prompt)\n",
139
+ " if st == -1:\n",
140
+ " return {\n",
141
+ " output1: out,\n",
142
+ " output2: None,\n",
143
+ " output3: None,\n",
144
+ " stat: 0\n",
145
+ " }\n",
146
+ " elif st == 0:\n",
147
+ " return {\n",
148
+ " output1: o1,\n",
149
+ " output2: out,\n",
150
+ " output3: None,\n",
151
+ " stat: 1\n",
152
+ " }\n",
153
+ " elif st == 1:\n",
154
+ " return {\n",
155
+ " output1: o1,\n",
156
+ " output2: o2,\n",
157
+ " output3: out,\n",
158
+ " stat: 2\n",
159
+ " }\n",
160
+ "\n",
161
+ "\n",
162
+ "with gr.Blocks() as app:\n",
163
+ " with gr.Row():\n",
164
+ " gr.Markdown(\n",
165
+ " \"\"\"\n",
166
+ " # NanoGPT - Cricket Commentary Generative AI\n",
167
+ " ### Give a prompt and see how it comes out with cricket commentary :)\n",
168
+ " \"\"\")\n",
169
+ "\n",
170
+ " with gr.Row(visible=True):\n",
171
+ " search_text = gr.Textbox(value=fn_query_on_load, placeholder='Enter prompt..', label='Enter Prompt')\n",
172
+ "\n",
173
+ " with gr.Row():\n",
174
+ " submit_btn = gr.Button(\"Submit\", variant='primary')\n",
175
+ " clear_btn = gr.ClearButton()\n",
176
+ " with gr.Row():\n",
177
+ " with gr.Column():\n",
178
+ " output1 = gr.Textbox(lines=10, interactive=False, label='Commentary Box')\n",
179
+ " output2 = gr.Textbox(lines=10, interactive=False, label='Commentary Box')\n",
180
+ " output3 = gr.Textbox(lines=10, interactive=False, label='Commentary Box')\n",
181
+ " stat = gr.State(value=-1)\n",
182
+ " \n",
183
+ "\n",
184
+ " def clear_data():\n",
185
+ " return {\n",
186
+ " output1: None,\n",
187
+ " output2: None,\n",
188
+ " output3: None,\n",
189
+ " search_text: None\n",
190
+ " }\n",
191
+ "\n",
192
+ " clear_btn.click(clear_data, None, [output1, output2, output3, search_text])\n",
193
+ "\n",
194
+ "\n",
195
+ " submit_btn.click(\n",
196
+ " fn_gen_comm,\n",
197
+ " [search_text, stat, output1, output2, output3],\n",
198
+ " [output1, output2, output3, stat]\n",
199
+ " )\n",
200
+ " \n",
201
+ " '''output1.change(\n",
202
+ " fn_gen_comm,\n",
203
+ " search_text,\n",
204
+ " [output1, output2, output3, stat]\n",
205
+ " )\n",
206
+ " \n",
207
+ " output2.change(\n",
208
+ " fn_gen_comm,\n",
209
+ " search_text,\n",
210
+ " [output1, output2, output3, stat]\n",
211
+ " )\n",
212
+ "\n",
213
+ " output3.change(\n",
214
+ " fn_gen_comm,\n",
215
+ " search_text,\n",
216
+ " [output1, output2, output3, stat]\n",
217
+ " )'''\n",
218
+ "\n",
219
+ "'''\n",
220
+ "Launch the app\n",
221
+ "'''\n",
222
+ "app.queue().launch()"
223
+ ]
224
+ },
225
+ {
226
+ "cell_type": "code",
227
+ "execution_count": null,
228
+ "metadata": {},
229
+ "outputs": [],
230
+ "source": []
231
+ }
232
+ ],
233
+ "metadata": {
234
+ "kernelspec": {
235
+ "display_name": "Python 3 (ipykernel)",
236
+ "language": "python",
237
+ "name": "python3"
238
+ },
239
+ "language_info": {
240
+ "codemirror_mode": {
241
+ "name": "ipython",
242
+ "version": 3
243
+ },
244
+ "file_extension": ".py",
245
+ "mimetype": "text/x-python",
246
+ "name": "python",
247
+ "nbconvert_exporter": "python",
248
+ "pygments_lexer": "ipython3",
249
+ "version": "3.8.5"
250
+ }
251
+ },
252
+ "nbformat": 4,
253
+ "nbformat_minor": 1
254
+ }