make sure <eos> is appended

lucidrains · Apr 2, 2021 · 4aade13 · 4aade13
1 parent b6ebceb
commit 4aade13
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/dalle_pytorch/simple_tokenizer.py b/dalle_pytorch/simple_tokenizer.py
@@ -122,12 +122,12 @@ def decode(self, tokens):
 
 tokenizer = SimpleTokenizer()
 
-def tokenize(texts, context_length = 256, add_start_and_end = False, truncate_text=False):
+def tokenize(texts, context_length = 256, add_start = False, add_end = True, truncate_text = False):
     if isinstance(texts, str):
         texts = [texts]
 
-    sot_tokens = [tokenizer.encoder["<|startoftext|>"]] if add_start_and_end else []
-    eot_tokens = [tokenizer.encoder["<|endoftext|>"]] if add_start_and_end else []
+    sot_tokens = [tokenizer.encoder["<|startoftext|>"]] if add_start else []
+    eot_tokens = [tokenizer.encoder["<|endoftext|>"]] if add_end else []
     all_tokens = [sot_tokens + tokenizer.encode(text) + eot_tokens for text in texts]
     result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
 

diff --git a/setup.py b/setup.py
@@ -4,7 +4,7 @@
   name = 'dalle-pytorch',
   packages = find_packages(),
   include_package_data = True,
-  version = '0.8.0',
+  version = '0.8.1',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',