Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import gradio as gr | |
| import tiktoken | |
| import math | |
| from dataclasses import dataclass | |
| class LayerNorm(nn.Module): | |
| def __init__(self, ndim, bias): | |
| super().__init__() | |
| self.weight = nn.Parameter(torch.ones(ndim)) | |
| self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None | |
| def forward(self, x): | |
| return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5) | |
| class CausalSelfAttention(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| assert config.n_embd % config.n_head == 0 | |
| self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias) | |
| self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) | |
| self.attn_dropout = nn.Dropout(config.dropout) | |
| self.resid_dropout = nn.Dropout(config.dropout) | |
| self.n_head, self.n_embd = config.n_head, config.n_embd | |
| self.flash = hasattr(F, 'scaled_dot_product_attention') | |
| if not self.flash: | |
| self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)) | |
| .view(1, 1, config.block_size, config.block_size)) | |
| def forward(self, x): | |
| B, T, C = x.size() | |
| q, k, v = self.c_attn(x).split(self.n_embd, dim=2) | |
| k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) | |
| q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) | |
| v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) | |
| y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True) | |
| y = y.transpose(1, 2).contiguous().view(B, T, C) | |
| return self.resid_dropout(self.c_proj(y)) | |
| class MLP(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias) | |
| self.gelu = nn.GELU() | |
| self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias) | |
| self.dropout = nn.Dropout(config.dropout) | |
| def forward(self, x): | |
| return self.dropout(self.c_proj(self.gelu(self.c_fc(x)))) | |
| class Block(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.ln1 = LayerNorm(config.n_embd, config.bias) | |
| self.attn = CausalSelfAttention(config) | |
| self.ln2 = LayerNorm(config.n_embd, config.bias) | |
| self.mlp = MLP(config) | |
| def forward(self, x): | |
| x = x + self.attn(self.ln1(x)) | |
| x = x + self.mlp(self.ln2(x)) | |
| return x | |
| class GPTConfig: | |
| block_size: int = 256 | |
| vocab_size: int = 50257 | |
| n_layer: int = 6 | |
| n_head: int = 6 | |
| n_embd: int = 384 | |
| dropout: float = 0.0 | |
| bias: bool = True | |
| class GPT(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.config = config | |
| self.transformer = nn.ModuleDict(dict( | |
| wte = nn.Embedding(config.vocab_size, config.n_embd), | |
| wpe = nn.Embedding(config.block_size, config.n_embd), | |
| drop = nn.Dropout(config.dropout), | |
| h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), | |
| ln_f = LayerNorm(config.n_embd, config.bias), | |
| )) | |
| self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) | |
| self.transformer.wte.weight = self.lm_head.weight | |
| self.apply(self._init_weights) | |
| def _init_weights(self, module): | |
| if isinstance(module, nn.Linear): | |
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| if module.bias is not None: | |
| torch.nn.init.zeros_(module.bias) | |
| elif isinstance(module, nn.Embedding): | |
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| def forward(self, idx, targets=None): | |
| device = idx.device | |
| b, t = idx.size() | |
| pos = torch.arange(0, t, dtype=torch.long, device=device) | |
| tok_emb = self.transformer.wte(idx) | |
| pos_emb = self.transformer.wpe(pos) | |
| x = self.transformer.drop(tok_emb + pos_emb) | |
| for block in self.transformer.h: | |
| x = block(x) | |
| x = self.transformer.ln_f(x) | |
| logits = self.lm_head(x[:, [-1], :]) | |
| return logits, None | |
| def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): | |
| for _ in range(max_new_tokens): | |
| idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:] | |
| logits, _ = self(idx_cond) | |
| logits = logits[:, -1, :] / temperature | |
| if top_k is not None: | |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) | |
| logits[logits < v[:, [-1]]] = -float('Inf') | |
| probs = F.softmax(logits, dim=-1) | |
| idx_next = torch.multinomial(probs, num_samples=1) | |
| idx = torch.cat((idx, idx_next), dim=1) | |
| return idx | |
| device = 'cpu' | |
| enc = tiktoken.get_encoding("gpt2") | |
| config = GPTConfig( | |
| vocab_size=50257, | |
| block_size=128, | |
| n_layer=6, | |
| n_head=6, | |
| n_embd=384, | |
| dropout=0.1, | |
| bias=True | |
| ) | |
| model = GPT(config) | |
| model.load_state_dict(torch.load('best_model_params.pt', map_location=device)) | |
| model.eval() | |
| model.to(device) | |
| def generate_story(prompt, max_new_tokens=100, temperature=0.8): | |
| start_ids = enc.encode_ordinary(prompt) | |
| x = torch.tensor(start_ids, dtype=torch.long, device=device)[None, ...] | |
| y = model.generate(x, max_new_tokens=int(max_new_tokens), temperature=temperature, top_k=200) | |
| return enc.decode(y[0].tolist()) | |
| article = """ | |
| <div style="text-align: center; margin-top: 20px;"> | |
| <p> | |
| This model was trained on the <strong>TinyStories dataset</strong> and uses a GPT-style architecture to generate new stories from a starting prompt. | |
| </p> | |
| <p> | |
| <strong>Created by:</strong> <a href="https://www.linkedin.com/in/harshitchawla4705/" target="_blank">Harshit Chawla</a> | |
| </p> | |
| <hr> | |
| <h4>Relevant Research Papers</h4> | |
| <p> | |
| <a href="https://arxiv.org/abs/2305.07759" target="_blank">TinyStories: How Small Can Language Models Be and Still Speak Coherent English?</a><br> | |
| <a href="https://d4mucfpksywv.cloudfront.net/better-language-models/language-models.pdf" target="_blank">Language Models are Unsupervised Multitask Learners (GPT-2)</a><br> | |
| <a href="https://arxiv.org/abs/1706.03762" target="_blank">Attention Is All You Need</a> | |
| </p> | |
| </div> | |
| """ | |
| iface = gr.Interface( | |
| fn=generate_story, | |
| inputs=[ | |
| gr.Textbox(lines=3, placeholder="E.g., Once upon a time, there was a little red fox..."), | |
| gr.Slider(minimum=20, maximum=500, step=10, value=150, label="Number of Words to Generate"), | |
| gr.Slider(minimum=0.1, maximum=1.5, step=0.1, value=0.8, label="Creativity (Temperature)") | |
| ], | |
| outputs=gr.Textbox(label="Generated Story", lines=20), | |
| title="📖 AI Story Generator", | |
| description="Enter a sentence to start a story and see what the AI writes next! Adjust the sliders to control the length and creativity of the story.", | |
| article=article, | |
| theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), | |
| allow_flagging="never", | |
| examples=[ | |
| ["Once upon a time, a curious cat named Whiskers", 150, 0.7], | |
| ["In a land filled with candy, a brave gingerbread man", 200, 0.9], | |
| ["A lonely robot sat on a hill, looking at the stars", 100, 0.5] | |
| ] | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch() |