Update app.py

app.py

@@ -17,6 +17,11 @@ def get_model():
 def generate_text(prompt, max_tokens=500, temperature=0.3, top_k=40):
     """Generate text based on the prompt."""
        
+    
+    
+    # Get cached model
+    model = get_model()
+    device = next(model.parameters()).device
     # Tokenize prompt with special token handling
     enc = tiktoken.get_encoding("gpt2")
     input_ids = torch.tensor(enc.encode(prompt, allowed_special={'<|endoftext|>'})).unsqueeze(0).to(device)
@@ -24,10 +29,6 @@ def generate_text(prompt, max_tokens=500, temperature=0.3, top_k=40):
     # Get end token id
     end_token = enc.encode('<|endoftext|>', allowed_special={'<|endoftext|>'})[0]
     
-    # Get cached model
-    model = get_model()
-    device = next(model.parameters()).device
-    
     with torch.no_grad():
         output_sequence = []
         progress_bar = st.progress(0)