Correct segment ID for XLNet single sequence

088fa7b7 · Lysandre · 073219b4 · 088fa7b7
Commit 088fa7b7 authored Jan 21, 2020 by Lysandre
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/transformers/tokenization_xlnet.py src/transformers/tokenization_xlnet.py +1 -1

No files found.
--- a/src/transformers/tokenization_xlnet.py
+++ b/src/transformers/tokenization_xlnet.py
@@ -240,7 +240,7 @@ class XLNetTokenizer(PreTrainedTokenizer):
        cls_segment_id = [2]

        if token_ids_1 is None:
-            return len(token_ids_0 + sep + cls) * [0]
+            return len(token_ids_0 + sep) * [0] + cls_segment_id
        return len(token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1] + cls_segment_id

    def save_vocabulary(self, save_directory):