Fix roberta overflowings
@@ -56,7 +56,41 @@ impl PostProcessor for RobertaProcessing {
offsets,
special_tokens,
attention_mask,
- encoding.take_overflowing(),
+ encoding
+ .take_overflowing()
+ .into_iter()
+ .map(|encoding| {
+ let ids = [&[self.cls.1], &encoding.get_ids()[..], &[self.sep.1]].concat();
+ let type_ids = [&[0], &encoding.get_type_ids()[..], &[0]].concat();
+ let tokens = [
+ &[self.cls.0.clone()],
+ &encoding.get_tokens()[..],
+ &[self.sep.0.clone()],
+ ]
+ .concat();
+ let words = [
+ &[0],
+ &encoding.get_words()[..],
+ &[encoding.get_words().last().map_or(0, |w| *w + 2)],
+ ]
+ .concat();
+ let offsets = [&[(0, 0)], &encoding.get_offsets()[..], &[(0, 0)]].concat();
+ let special_tokens =
+ [&[1u32], &vec![0; encoding.get_ids().len()][..], &[1]].concat();
+ let attention_mask = vec![1; ids.len()];
+
+ Encoding::new(
+ ids,
+ type_ids,
+ tokens,
+ words,
+ offsets,
+ special_tokens,
+ attention_mask,
+ vec![],
+ )
+ })
+ .collect(),
);
if let Some(mut encoding) = pair_encoding {
@@ -87,7 +121,43 @@ impl PostProcessor for RobertaProcessing {
pair_offsets,
pair_special_tokens,
pair_attention_mask,
- encoding.take_overflowing(),
+ encoding
+ .take_overflowing()
+ .into_iter()
+ .map(|encoding| {
+ let pair_ids =
+ [&[self.sep.1], &encoding.get_ids()[..], &[self.sep.1]].concat();
+ let pair_type_ids = vec![0; encoding.get_ids().len() + 2];
+ let pair_tokens = [
+ &[self.sep.0.clone()],
+ &encoding.get_tokens()[..],
+ &[self.sep.0.clone()],
+ ]
+ .concat();
+ let pair_words = [
+ &[0],
+ &encoding.get_words()[..],
+ &[encoding.get_words().last().map_or(0, |w| *w + 2)],
+ ]
+ .concat();
+ let pair_offsets =
+ [&[(0, 0)], &encoding.get_offsets()[..], &[(0, 0)]].concat();
+ let pair_special_tokens =
+ [&[1], &vec![0u32; encoding.get_type_ids().len()][..], &[1]].concat();
+ let pair_attention_mask = vec![1; pair_ids.len()];
+
+ Encoding::new(
+ pair_ids,
+ pair_type_ids,
+ pair_tokens,
+ pair_words,
+ pair_offsets,
+ pair_special_tokens,
+ pair_attention_mask,
+ vec![],
+ )
+ })
+ .collect(),
);
new_encoding.merge_with(new_pair_encoding, false);
