Update app.py

app.py

@@ -129,8 +129,6 @@ def load_model(model_path):
     model.to(device)
     return model
 
-# Don't load the model here
-# model = load_model('gpt_model.pth')
 enc = tiktoken.get_encoding('gpt2')
 
 # Update the generate_text function
@@ -166,6 +164,14 @@ async def generate_text(prompt, max_length=432, temperature=0.8, top_k=40):
     if len(generated) == max_length:
         yield "... (output truncated due to length)"
 
+# Add the gradio_generate function
+@spaces.GPU(duration=60)
+async def gradio_generate(prompt, max_length, temperature, top_k):
+    output = ""
+    async for token in generate_text(prompt, max_length, temperature, top_k):
+        output += token
+        yield output
+
 
 # # Your existing imports and model code here...