SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Sequence special token handling for BERT and RoBERTa

Browse Source
This commit is contained in:
Lysandre
2019-08-20 14:15:28 -04:00
parent a690edab17
commit 2d042274ac
1 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
examples/run_lm_finetuning.py
View File

@@ -71,9 +71,15 @@ class TextDataset(Dataset):
text = f.read()
tokenized_text = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(text))
tokenized_text = tokenizer.add_special_tokens_single_sentence(tokenized_text)
while len(tokenized_text) >= block_size: # Truncate in block of block_size
self.examples.append(tokenized_text[:block_size])
tokenized_text = tokenized_text[block_size:]
if isinstance(tokenizer, (BertTokenizer, RobertaTokenizer)):
self.examples.append(tokenizer.add_special_tokens_single_sentence(tokenized_text[:block_size - 2]))
tokenized_text = tokenized_text[block_size - 2:]
else:
self.examples.append(tokenized_text[:block_size])
tokenized_text = tokenized_text[block_size:]
# Note that we are loosing the last truncated example here for the sake of simplicity (no padding)
# If your dataset is small, first you should loook for a bigger one :-) and second you
# can change this behavior by adding (model specific) padding.