Add get_*_segment_ids functions
@@ -214,6 +214,50 @@ def is_class_token(text: str) -> bool:
return text == "[CLS]" or text == "<cls>"
+def get_bert_segment_ids(length1: int, length2: int, *, xp=numpy) -> Array:
+ """Get an array of segment IDs in BERT's format, for an input with one or
+ two segments (set length2=0 for one segment). The lengths should be just the
+ wordpiece lengths, not including the SEP and CLS tokens.
+
+ According to the HF glue_utils.py module, the convention for BERT is:
+
+ (a) For sequence pairs:
+ tokens: [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
+ type_ids: 0 0 0 0 0 0 0 0 1 1 1 1 1 1
+
+ (b) For single sequences:
+ tokens: [CLS] the dog is hairy . [SEP]
+ type_ids: 0 0 0 0 0 0 0
+ """
+ if length2:
+ ids = [0] * (length1+2) + [1] * (length2 + 1)
+ else:
+ ids = [0] * (length1+2)
+ return xp.array(ids, dtype=xp.int_)
+
+
+def get_xlnet_segment_ids(length1: int, length2: int, *, xp=numpy) -> Array:
+ """Get an array of segment IDs in XLNet's format, for an input with one or
+ two segments (set length2=0 for one segment). The lengths should be just the
+ wordpiece lengths, not including the SEP and CLS tokens.
+
+ According to the XLNet code classifer_utils.py module, the convention is:
+
+ (a) For sequence pairs:
+ tokens: is this jack ##son ##ville ? <sep> no it is not . <sep> <cls>
+ type_ids: 0 0 0 0 0 0 0 1 1 1 1 1 1 2
+
+ (b) For single sequences:
+ tokens: the dog is hairy . <sep> <cls>
+ type_ids: 0 0 0 0 0 0 2
+ """
+ if length2:
+ ids = [0] * length1 + [0] + [1] * length2 + [1, 2]
+ else:
+ ids = [0] * length1 + [0, 2]
+ return xp.array(ids, dtype=xp.int_)
+
+
def get_sents(doc: Union[Span, Doc]) -> List[Span]:
if doc.is_sentenced:
return list(doc.sents)
