Create HF Spaces app

.gitattributes

@@ -0,0 +1 @@
+checkpoints/model.pt filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,5 +1,5 @@
 # Large files
-checkpoints/
+checkpoints/candidates
 datasets/
 
 # Training Tensorboard runs

app.py

@@ -0,0 +1,33 @@
+import gradio
+import gpt
+
+# Intended for HuggingFace Spaces for GPT
+config = gpt.GPTConfig("gpt_config.json")
+model = gpt.GPT(
+    vocab_size=config.vocab_size,
+    context_size=config.context_size,
+    embedding_dim=config.embedding_dim
+)
+
+# Use GPU if enabled
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# Load checkpoint
+model, _, _ = load_checkpoint(model, None, config.restore_path)
+
+demo = gr.Interface(
+    fn=lambda *args: generate(model, config, *args),
+    inputs=[
+        gr.Textbox(lines=2, placeholder="Prompt here..."),
+        gr.Number(precision=0, value=256),
+        gr.Number(value=0.8),
+        gr.Slider(maximum=128, value=10),
+        gr.Slider(maximum=1, value=1)
+    ],
+    outputs="text",
+    title="Shakespeare-GPT",
+    description="Putting theater kids out of their nonexistent jobs since 2023"
+)
+
+demo.launch()

gpt.py

@@ -281,7 +281,8 @@ def load_checkpoint(model, optimizer, path):
     """
     checkpoint = torch.load(path)
     model.load_state_dict(checkpoint['model_state_dict'])
-    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    if optimizer is not None:
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
     return (model, optimizer, checkpoint['steps'])