|endoftext| token handled

app.py

@@ -31,7 +31,8 @@ model = load_model_from_hf()
 model.train(False)
 
 def generate_text(prompt, max_length=100, num_samples=1, temperature=0.8):
-    enc = tiktoken.get_encoding('gpt2')
+    # Initialize encoder with allowed special tokens
+    enc = tiktoken.get_encoding('gpt2', allowed_special={'<|endoftext|>'})
     tokens = enc.encode(prompt)
     tokens = torch.tensor(tokens, dtype=torch.long)
     tokens = tokens.unsqueeze(0).repeat(num_samples, 1)
@@ -54,7 +55,7 @@ def generate_text(prompt, max_length=100, num_samples=1, temperature=0.8):
             tokens = torch.cat((tokens, next_token), dim=1)
             
             # Check for end of text token
-            if next_token.item() == enc.encode('<|endoftext|>')[0]:
+            if next_token.item() == enc.encode('<|endoftext|>', allowed_special={'<|endoftext|>'})[0]:
                 break
     
     generated_texts = []