""" Shakespeare Text Generator - Hugging Face Gradio App ===================================================== A GPT model trained on Shakespeare's works to generate text in Shakespearean style. """ import gradio as gr import torch import torch.nn as nn from torch.nn import functional as F import math from dataclasses import dataclass import tiktoken # ============================================================================ # MODEL ARCHITECTURE (Same as training) # ============================================================================ 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) self.c_proj = nn.Linear(config.n_embd, config.n_embd) self.c_proj.NANGPT_SCALE_INIT = 1 self.n_head = config.n_head self.n_embd = config.n_embd 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() qkv = self.c_attn(x) q, k, v = qkv.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) att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf')) att = F.softmax(att, dim=-1) y = att @ v y = y.transpose(1, 2).contiguous().view(B, T, C) y = self.c_proj(y) return y class MLP(nn.Module): def __init__(self, config): super().__init__() self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd) self.gelu = nn.GELU(approximate='tanh') self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd) self.c_proj.NANOGPT_SCALE_INIT = 1 def forward(self, x): x = self.c_fc(x) x = self.gelu(x) x = self.c_proj(x) return x class Block(nn.Module): def __init__(self, config): super().__init__() self.ln_1 = nn.LayerNorm(config.n_embd) self.attn = CausalSelfAttention(config) self.ln_2 = nn.LayerNorm(config.n_embd) self.mlp = MLP(config) def forward(self, x): x = x + self.attn(self.ln_1(x)) x = x + self.mlp(self.ln_2(x)) return x @dataclass class GPTConfig: block_size: int = 1024 vocab_size: int = 50257 n_layer: int = 12 n_head: int = 12 n_embd: int = 768 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), h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), ln_f = nn.LayerNorm(config.n_embd), )) 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): std = 0.02 if hasattr(module, 'NANGPT_SCALE_INIT'): std *= (2 * self.config.n_layer) ** -0.5 torch.nn.init.normal_(module.weight, mean=0.0, std=std) 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): B, T = idx.size() assert T <= self.config.block_size pos = torch.arange(0, T, dtype=torch.long, device=idx.device) pos_emb = self.transformer.wpe(pos) tok_emb = self.transformer.wte(idx) x = tok_emb + pos_emb for block in self.transformer.h: x = block(x) x = self.transformer.ln_f(x) logits = self.lm_head(x) loss = None if targets is not None: loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) return logits, loss # ============================================================================ # LOAD MODEL # ============================================================================ print("Loading model...") device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"Using device: {device}") model = GPT(GPTConfig()) # Load your trained checkpoint here checkpoint = torch.load('shakespeare_gpt_fp16.pt', map_location=device) model.load_state_dict({k: v.float() for k, v in checkpoint.items()}) model.to(device) model.eval() enc = tiktoken.get_encoding('gpt2') # ============================================================================ # GENERATION FUNCTION # ============================================================================ def generate_shakespeare(prompt, max_length=100, temperature=0.8, top_k=50, num_samples=1): """ Generate Shakespeare-style text from a prompt. Args: prompt: Starting text max_length: Maximum number of tokens to generate temperature: Sampling temperature (higher = more random) top_k: Number of top tokens to sample from num_samples: Number of different samples to generate """ if not prompt.strip(): return "Please enter a prompt to generate text." try: # Encode the prompt tokens = enc.encode(prompt) if len(tokens) == 0: return "Invalid prompt. Please try again." outputs = [] for _ in range(num_samples): x = torch.tensor(tokens, dtype=torch.long, device=device).unsqueeze(0) with torch.no_grad(): for _ in range(max_length): # Forward pass logits = model(x)[0] logits = logits[:, -1, :] / temperature # Top-k sampling probs = F.softmax(logits, dim=-1) topk_probs, topk_indices = torch.topk(probs, min(top_k, probs.size(-1)), dim=-1) ix = torch.multinomial(topk_probs, 1) xcol = torch.gather(topk_indices, -1, ix) x = torch.cat((x, xcol), dim=1) # Stop if we exceed block size if x.size(1) >= model.config.block_size: break # Decode the output output_tokens = x[0].tolist() generated_text = enc.decode(output_tokens) outputs.append(generated_text) # Return all samples separated by dividers if num_samples == 1: return outputs[0] else: return "\n\n" + "="*60 + "\n\n".join(outputs) except Exception as e: return f"Error generating text: {str(e)}" # ============================================================================ # GRADIO INTERFACE # ============================================================================ # Create the interface with gr.Blocks() as demo: gr.Markdown( """ # 🎭 Shakespeare Text Generator Generate text in the style of William Shakespeare using a GPT model trained on his complete works. Enter a prompt and watch the Bard's AI apprentice continue the story! **Model Details**: GPT-2 124M architecture trained on Shakespeare's plays and sonnets (Loss: 0.095) """ ) with gr.Row(): with gr.Column(scale=1): prompt_input = gr.Textbox( label="📝 Enter Your Prompt", placeholder="To be or not to be...", lines=4, value="To be or not to be" ) with gr.Accordion("⚙️ Advanced Settings", open=False): max_length_slider = gr.Slider( minimum=20, maximum=300, value=100, step=10, label="Max Length (tokens)", info="Maximum number of tokens to generate" ) temperature_slider = gr.Slider( minimum=0.1, maximum=1.5, value=0.8, step=0.1, label="Temperature", info="Higher = more creative, Lower = more focused" ) top_k_slider = gr.Slider( minimum=10, maximum=100, value=50, step=10, label="Top-K", info="Number of top tokens to sample from" ) num_samples_slider = gr.Slider( minimum=1, maximum=3, value=1, step=1, label="Number of Samples", info="Generate multiple variations" ) generate_btn = gr.Button("🎨 Generate", variant="primary", size="lg") with gr.Column(scale=1): output_text = gr.Textbox( label="📜 Generated Text", lines=15 ) # Examples gr.Markdown("### 📚 Try These Examples:") gr.Examples( examples=[ ["To be or not to be", 100, 0.8, 50, 1], ["What's in a name?", 120, 0.7, 40, 1], ["All the world's a stage", 150, 0.9, 50, 1], ["Romeo, Romeo, wherefore art thou", 100, 0.8, 50, 1], ["Friends, Romans, countrymen", 130, 0.75, 45, 1], ["Now is the winter of our discontent", 110, 0.85, 50, 1], ], inputs=[prompt_input, max_length_slider, temperature_slider, top_k_slider, num_samples_slider], outputs=output_text, fn=generate_shakespeare, cache_examples=False ) # Connect the button generate_btn.click( fn=generate_shakespeare, inputs=[prompt_input, max_length_slider, temperature_slider, top_k_slider, num_samples_slider], outputs=output_text ) gr.Markdown( """ --- ### 💡 Tips for Best Results: - Start with famous Shakespeare quotes for coherent continuations - Use **lower temperature** (0.5-0.7) for more focused, coherent text - Use **higher temperature** (0.9-1.2) for more creative, diverse outputs - Adjust **Top-K** to control vocabulary diversity - Try generating multiple samples to see different variations ### ⚠️ Note: This model was trained on Shakespeare's works and will generate text in Early Modern English style. Results may vary based on the prompt and parameters. """ ) # ============================================================================ # LAUNCH # ============================================================================ if __name__ == "__main__": demo.launch(server_port=7860, share=True)