|endoftext| token handled

app.py

@@ -31,8 +31,9 @@ model = load_model_from_hf()
 model.train(False)
 
 def generate_text(prompt, max_length=100, num_samples=1, temperature=0.8):
-    # Initialize encoder with allowed special tokens
-    enc = tiktoken.get_encoding('gpt2', allowed_special={'<|endoftext|>'})
+    enc = tiktoken.get_encoding('gpt2')
+    # Modify encoding behavior to allow special tokens
+    enc._special_tokens.add('<|endoftext|>')
     tokens = enc.encode(prompt)
     tokens = torch.tensor(tokens, dtype=torch.long)
     tokens = tokens.unsqueeze(0).repeat(num_samples, 1)
@@ -55,7 +56,8 @@ 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|>', allowed_special={'<|endoftext|>'})[0]:
+            endoftext_token = enc.encode('<|endoftext|>')[0]
+            if next_token.item() == endoftext_token:
                 break
     
     generated_texts = []