Create app.py

app.py

@@ -0,0 +1,366 @@
+"""
+Shakespeare Text Generator - Hugging Face Gradio App
+Trained GPT-2 model (124M params)
+"""
+
+import gradio as gr
+import torch
+import tiktoken
+import os
+from dataclasses import dataclass
+
+
+# GPT Model Architecture
+@dataclass
+class GPTConfig:
+    block_size: int = 1024
+    vocab_size: int = 50257
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.0
+    bias: bool = True
+
+
+import torch.nn as nn
+from torch.nn import functional as F
+import math
+
+
+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 = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
+                            .view(1, 1, config.block_size, config.block_size))
+        self.c_proj.NANOGPT_SCALE_INIT = 1
+    
+    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)
+        att = self.attn_dropout(att)
+        y = att @ v
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        y = self.resid_dropout(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, bias=config.bias)
+        self.gelu = nn.GELU(approximate='tanh')
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+        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)
+        x = self.dropout(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
+
+
+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=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, 'NANOGPT_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):
+        device = idx.device
+        b, t = idx.size()
+        assert t <= self.config.block_size
+        pos = torch.arange(0, t, dtype=torch.long, device=device)
+        pos_emb = self.transformer.wpe(pos)
+        tok_emb = self.transformer.wte(idx)
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+        if targets is not None:
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            logits = self.lm_head(x[:, [-1], :])
+            loss = None
+        return logits, loss
+    
+    @torch.no_grad()
+    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
+
+
+# Load model
+print("Loading model...")
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+config = GPTConfig()
+model = GPT(config)
+
+# Load checkpoint
+checkpoint_path = "model_quantized.pt"
+if os.path.exists(checkpoint_path):
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    print(f"✓ Loaded quantized model from {checkpoint_path}")
+    print(f"  Training loss: {checkpoint.get('loss', 'N/A')}")
+    print(f"  Model size: 330MB (FP16 quantized)")
+else:
+    print("⚠️ Checkpoint not found. Please upload 'model_quantized.pt'")
+
+model.to(device)
+model.eval()
+print(f"✓ Model ready on {device}")
+
+# Tokenizer
+enc = tiktoken.get_encoding('gpt2')
+
+
+# ---- Derived Stats (dynamic, for UI) ----
+try:
+    model_params = sum(p.numel() for p in model.parameters())
+    model_params_m = model_params / 1e6
+except Exception:
+    model_params = None
+    model_params_m = None
+
+training_loss = None
+training_step = None
+if 'checkpoint' in locals():
+    training_loss = checkpoint.get('loss', None)
+    training_step = checkpoint.get('step', None)
+
+def build_stats_md() -> str:
+    params_line = f"- **Parameters**: {model_params:,} ({model_params_m:.0f}M)" if model_params is not None else "- **Parameters**: 124M"
+    loss_line = f"- **Training Loss**: {training_loss:.6f}" if isinstance(training_loss, (float, int)) else "- **Training Loss**: N/A"
+    step_line = f"- **Training Step**: {training_step}" if training_step is not None else "- **Training Step**: N/A"
+    return f"""
+### 📊 Model Details
+{params_line}
+- **Architecture**: GPT-2 (Decoder-only Transformer)
+{loss_line}
+{step_line}
+- **Model Format**: FP16 quantized (≈330MB)
+- **Device**: {device.upper()}
+""".strip()
+
+
+def generate_text(prompt, max_tokens=100, temperature=0.8, top_k=50):
+    """Generate text from a prompt"""
+    
+    if not prompt:
+        return "⚠️ Please enter a prompt!"
+    
+    try:
+        # Tokenize
+        tokens = enc.encode(prompt)
+        tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(device)
+        
+        # Generate
+        with torch.no_grad():
+            generated = model.generate(
+                tokens,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                top_k=top_k if (top_k and int(top_k) > 0) else None
+            )
+        
+        # Decode
+        generated_text = enc.decode(generated[0].tolist())
+        
+        return generated_text
+    
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+
+# Example prompts
+examples = [
+    ["First Citizen:", 150, 0.8, 50],
+    ["ROMEO:", 150, 0.8, 50],
+    ["To be, or not to be,", 200, 0.7, 40],
+    ["What light through yonder window breaks?", 150, 0.8, 50],
+    ["Friends, Romans, countrymen,", 150, 0.8, 50],
+]
+
+
+# Gradio Interface with Teal Theme
+with gr.Blocks(
+    title="Shakespeare Text Generator",
+    theme=gr.themes.Soft(
+        primary_hue="teal",
+        secondary_hue="cyan",
+        neutral_hue="slate"
+    ),
+    css="""
+    .gradio-container {
+        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+    }
+    .gr-button-primary {
+        background: linear-gradient(135deg, #14b8a6 0%, #0d9488 100%) !important;
+        border: none !important;
+        color: white !important;
+        font-weight: 600 !important;
+    }
+    .gr-button-primary:hover {
+        background: linear-gradient(135deg, #0d9488 0%, #0f766e 100%) !important;
+        transform: translateY(-1px);
+        box-shadow: 0 4px 12px rgba(20, 184, 166, 0.3) !important;
+    }
+    h1 {
+        color: #0f766e !important;
+        text-align: center;
+    }
+    .badge {
+        display: inline-block;
+        padding: 6px 10px;
+        margin: 4px 6px 0 0;
+        border-radius: 8px;
+        background: #ecfeff;
+        color: #0f766e;
+        font-size: 12px;
+        border: 1px solid #ccfbf1;
+    }
+    """
+) as demo:
+    gr.Markdown(f"""
+    # 🎭 Shakespeare Text Generator
+    
+    <div>
+      <span class="badge">Architecture: GPT-2</span>
+      <span class="badge">Device: {device.upper()}</span>
+      <span class="badge">Quantized: FP16</span>
+      <span class="badge">Params: {int(model_params_m):d}M</span>
+    </div>
+    
+    Enter a Shakespearean prompt and watch the AI continue the text!
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            prompt_input = gr.Textbox(
+                label="Prompt",
+                placeholder="Enter a Shakespearean prompt (e.g., 'First Citizen:', 'ROMEO:', 'To be, or not to be,')",
+                lines=3
+            )
+            
+            with gr.Accordion("Advanced Settings", open=False):
+                with gr.Row():
+                    max_tokens = gr.Slider(
+                        minimum=50,
+                        maximum=600,
+                        value=150,
+                        step=10,
+                        label="Max Tokens"
+                    )
+                    temperature = gr.Slider(
+                        minimum=0.5,
+                        maximum=1.5,
+                        value=0.8,
+                        step=0.1,
+                        label="Temperature (creativity)"
+                    )
+                    top_k = gr.Slider(
+                        minimum=0,
+                        maximum=100,
+                        value=50,
+                        step=10,
+                        label="Top-K (diversity) (0 disables)"
+                    )
+            
+            generate_btn = gr.Button("✨ Generate Shakespeare", variant="primary", size="lg")
+        
+        with gr.Column(scale=2):
+            output_text = gr.Textbox(
+                label="Generated Text",
+                lines=15,
+                show_copy_button=True
+            )
+
+            gr.Markdown(build_stats_md())
+    
+    gr.Markdown("""
+    ### 💡 Tips:
+    - **Temperature**: Lower (0.5-0.7) = more focused, Higher (0.9-1.2) = more creative
+    - **Top-K**: Controls vocabulary diversity (40-60 recommended)
+    - **Prompts**: Try character names (ROMEO:, JULIET:) or famous phrases
+    """)
+    
+    gr.Examples(
+        examples=examples,
+        inputs=[prompt_input, max_tokens, temperature, top_k],
+        label="Example Prompts"
+    )
+    
+    gr.Markdown(build_stats_md())
+    
+    # Connect button
+    generate_btn.click(
+        fn=generate_text,
+        inputs=[prompt_input, max_tokens, temperature, top_k],
+        outputs=output_text
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()
+