Upload app.py

app.py

@@ -98,23 +98,17 @@ class GPT(nn.Module):
             loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
         return logits, loss
 
-# Load model
+# Load model and tokenizer
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-print(f"Loading model on {device}...")
 config = GPTConfig()
 model = GPT(config)
 model_path = os.path.join("models", "best_model.pt")
-#checkpoint = torch.load(model_path, map_location=device, weights_only=False)
-#model.load_state_dict(checkpoint['model_state_dict'])
-
-model.load_state_dict(torch.load(model_path, map_location=device, weights_only=False))
+model.load_state_dict(torch.load(model_path, map_location=device))
 model.to(device)
 model.eval()
-
 enc = tiktoken.get_encoding('gpt2')
 
-print(f"✅ Model loaded!")
-
+print(f"✅ Model loaded on {device}!")
 
 def generate(prompt: str, max_new_tokens: int = 30, top_k: int = 50, temperature: float = 1.0):
     tokens = enc.encode(prompt)
@@ -134,15 +128,48 @@ def generate(prompt: str, max_new_tokens: int = 30, top_k: int = 50, temperature
     out_tokens = x[0].tolist()
     return enc.decode(out_tokens)
 
+# Example prompts for dropdown
+example_prompts = [
+    "To be, or not to be, that is the question:",
+    "O Romeo, Romeo! wherefore art thou Romeo?",
+    "Once more unto the breach, dear friends, once more;",
+    "All the world's a stage,",
+    "The lady doth protest too much, methinks."
+]
+
 with gr.Blocks() as demo:
-    gr.Markdown("# GPT2-Space")
+    gr.Markdown("# GPT-2 (124M) Shakespeare Text Generator")
+    gr.Markdown(
+        "GPT-2 (124M) model trained from scratch on Shakespeare's works. "
+        "Start with a prompt and generate Shakespearean-style text!"
+    )
+
     with gr.Row():
         inp = gr.Textbox(lines=3, placeholder="Enter prompt here...", label="Prompt")
-        out = gr.Textbox(lines=10, label="Generated")
+        out = gr.Textbox(lines=10, label="Generated Text")
+
     with gr.Row():
         max_tokens = gr.Slider(1, 200, value=30, step=1, label="Max new tokens")
         topk = gr.Slider(1, 200, value=50, step=1, label="Top-k")
         temp = gr.Slider(0.01, 2.0, value=1.0, step=0.01, label="Temperature")
+
+    with gr.Row():
+        example_dropdown = gr.Dropdown(
+            choices=example_prompts,
+            label="Choose example prompt",
+            interactive=True
+        )
+        clear_btn = gr.Button("Clear output")
+
+    def use_example(prompt):
+        return prompt
+
+    def clear_output():
+        return ""
+
+    example_dropdown.change(fn=use_example, inputs=example_dropdown, outputs=inp)
+    clear_btn.click(fn=clear_output, inputs=[], outputs=out)
+
     btn = gr.Button("Generate")
     btn.click(fn=generate, inputs=[inp, max_tokens, topk, temp], outputs=out)