Feat: Zero GPU for inference

app.py

@@ -6,6 +6,7 @@ from dataclasses import dataclass
 import torch.nn as nn
 import math
 import inspect
+import spaces
 
 # Configuration class (same as in training)
 @dataclass
@@ -129,7 +130,8 @@ def load_model():
     return model, device
 
 # Text generation function
-def generate_text(prompt, num_tokens, model, device, temperature=0.8):
+@spaces.gpu(enable_queue=True)
+def generate_text(prompt, num_tokens, temperature=0.8):
     enc = tiktoken.get_encoding('gpt2')
     x = torch.tensor([enc.encode(prompt)], dtype=torch.long, device=device)
 
@@ -148,13 +150,9 @@ def generate_text(prompt, num_tokens, model, device, temperature=0.8):
 # Load the model globally
 model, device = load_model()
 
-# Gradio interface
-def gradio_interface(prompt, num_tokens, temperature):
-    return generate_text(prompt, num_tokens, model, device, temperature)
-
 # Create the Gradio interface
-iface = gr.Interface(
-    fn=gradio_interface,
+demo = gr.Interface(
+    fn=generate_text,
     inputs=[
         gr.Textbox(label="Enter your prompt", value="Once upon a time"),
         gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Number of tokens to generate"),
@@ -163,7 +161,9 @@ iface = gr.Interface(
     outputs=gr.Textbox(label="Generated Text"),
     title="NanoGPT Text Generator",
     description="Generate Shakespeare-style text using a trained NanoGPT model",
+    allow_flagging="never",
+    cache_examples=True
 )
 
 if __name__ == "__main__":
-    iface.launch() 
+    demo.launch()