Update app.py

app.py

@@ -17,8 +17,14 @@ def get_model():
 def generate_text(prompt, max_tokens=500, temperature=0.3, top_k=40):
     """Generate text based on the prompt."""
     # Encode the prompt
-    enc = tiktoken.get_encoding('gpt2')
-    input_ids = torch.tensor(enc.encode(prompt)).unsqueeze(0)
+    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)
+    
+    # Get end token id
+    end_token = enc.encode('<|endoftext|>', allowed_special={'<|endoftext|>'})[0]
     
     # Get cached model
     model = get_model()
@@ -31,7 +37,7 @@ def generate_text(prompt, max_tokens=500, temperature=0.3, top_k=40):
             progress_bar.progress(i / max_tokens)
             
             # Get predictions
-            outputs = model(input_ids)
+            logits,_ = model(input_ids)
             logits = outputs.logits[:, -1, :] / temperature
             
             # Apply top-k filtering