We want to fine tune BERT with Next Sentence Prediction (NSP) objective and we have a list of files which contains the conversations. To prepare the training dataset for the fine tuning, currently we read through all the files, load all conversation sentences into memory, create positive examples for adjacent sentences A and B, like  [CLS] A [SEP] B [SEP], and create negative examples by randomly sample two sentences A and B in all conversation sentences.
The Current logic is similar with:
  

      github.com
  

  
    huggingface/transformers/blob/master/src/transformers/data/datasets/language_modeling.py#L196 1



      
      
                  for i in range(n):
                      self.examples[i]["chinese_ref"] = torch.tensor(ref[i], dtype=torch.long)
          
          
    def __len__(self):
                  return len(self.examples)
          
          
    def __getitem__(self, i) -> Dict[str, torch.tensor]:
                  return self.examples[i]
          
          

          
class LineByLineWithSOPTextDataset(Dataset):
              """
              Dataset for sentence order prediction task, prepare sentence pairs for SOP task
              """
          
          
    def __init__(self, tokenizer: PreTrainedTokenizer, file_dir: str, block_size: int):
                  warnings.warn(
                      DEPRECATION_WARNING.format(
                          "https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/run_mlm.py"
                      ),
                      FutureWarning,
      
    
  


  

  
    
    
  

  


However, this is not memory efficient because it loads all sentences into memory and now we have lots of sentences which cannot fit into memory any more.
Any suggestions to create the dataset for NSP more memory efficiently? The load_dataset APIs look promising, but didn’t figure out how to process the input files to randomly sample sentences for the negative examples.
Thanks