Fixing sequence clipping bug in tokenizer (#46)

OpenBioML · Dec 21, 2023 · c92083b · c92083b
1 parent 0e48cff
commit c92083b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/protein_lm/tokenizer/tokenizer.py b/protein_lm/tokenizer/tokenizer.py
@@ -51,6 +51,8 @@ def batch_encode(
         output = []
         if max_sequence_length is None and return_tensors:
             max_sequence_length = max([len(sequence) for sequence in sequences])
+            if add_special_tokens:
+                max_sequence_length += 2
         if max_sequence_length is not None:
             sequences = [
                 sequence[:(max_sequence_length - 2) if add_special_tokens else max_sequence_length]