Fix to allow masked token after 512th token

Sequences longer than 510 are now truncated around the masked token for xlm-roberta-base, regardless of mask location.

app.py

@@ -31,8 +31,35 @@ xlmr_tokenizer = AutoTokenizer.from_pretrained('xlm-roberta-base', max_length=51
 xlmr_p = pipeline("fill-mask", model=model, tokenizer=tokenizer)
 
 def xlmr_base_fn(text):
- text = ' '.join(text.split()[:500])
- preds = xlmr_p(text)
+ # Find our masked token
+ tokens = xlmr_tokenizer.tokenize(text)
+ mask_token_idx = [i for i, x in enumerate(tokens) if xlmr_tokenizer.mask_token in x][0]
+ 
+ max_len = tokenizer.model_max_length
+ max_len = max_len-2 if max_len % 512 == 0 and max_len < 4096 else 510
+
+ # Smart truncation for long sequences
+ if not len(tokens) < max_len:
+
+ # Find left and right bounds for truncated sequences
+ lbound = max(0, mask_token_idx-(max_len//2))
+ rbound = min(len(tokens), mask_token_idx+(max_len//2))
+
+ # If we hit an edge, expand sequence in the other direction
+ if lbound == 0 and rbound != len(tokens)-1:
+ rbound = min(len(tokens), max_len)
+ elif rbound == len(tokens) and lbound != 0:
+ lbound = max(0, len(tokens)-max_len)
+
+ # Apply truncation and rejoin tokens to form new text
+ truncated_text = ''.join(tokens[lbound:rbound])
+
+ # Handle lowbar from xlmr tokenizer
+ truncated_text = ''.join([x if ord(x) != 9601 else ' ' for x in result])
+ else:
+ truncated_text = text
+ 
+ preds = xlmr_p(truncated_text)
  pred_dict = {}
  for pred in preds:
  pred_dict[pred['token_str']] = pred['score']