Update app.py

app.py

@@ -3,19 +3,16 @@ import torch.nn as nn
 from torch.nn import functional as F
 import tiktoken
 import gradio as gr
-import asyncio
 
-# Try to import spaces, use a dummy decorator if not available
 try:
     import spaces
     use_spaces_gpu = True
 except ImportError:
     use_spaces_gpu = False
-    # Dummy decorator in case spaces is not available
     def dummy_gpu_decorator(func):
         return func
     spaces = type('', (), {'GPU': dummy_gpu_decorator})()
-
+    
 # Define the GPTConfig class
 class GPTConfig:
     def __init__(self):
@@ -131,10 +128,10 @@ def load_model(model_path):
 
 enc = tiktoken.get_encoding('gpt2')
 
+
 # Update the generate_text function
-@spaces.GPU(duration=60)
-async def generate_text(prompt, max_length=432, temperature=0.8, top_k=40):
-    # Load the model inside the GPU-decorated function
+@spaces.GPU
+def generate_text(prompt, max_length=432, temperature=0.8, top_k=40):
     model = load_model('gpt_model.pth')
     device = next(model.parameters()).device
     input_ids = torch.tensor(enc.encode(prompt)).unsqueeze(0).to(device)
@@ -153,24 +150,15 @@ async def generate_text(prompt, max_length=432, temperature=0.8, top_k=40):
             input_ids = torch.cat([input_ids, next_token], dim=-1)
             generated.append(next_token.item())
             
-            next_token_str = enc.decode([next_token.item()])
-            yield next_token_str
-            
             if next_token.item() == enc.encode('\n')[0] and len(generated) > 100:
                 break
-            
-            await asyncio.sleep(0.02)
 
-    if len(generated) == max_length:
-        yield "... (output truncated due to length)"
+    return enc.decode(generated)
 
 # 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
+@spaces.GPU
+def gradio_generate(prompt, max_length, temperature, top_k):
+    return generate_text(prompt, max_length, temperature, top_k)
 
 
 # # Your existing imports and model code here...