--- a
+++ b/dataLoader/MLM.py
@@ -0,0 +1,53 @@
+from torch.utils.data.dataset import Dataset
+import numpy as np
+from dataLoader.utils import seq_padding,position_idx,index_seg,random_mask
+import torch
+
+
+class MLMLoader(Dataset):
+    def __init__(self, dataframe, token2idx, age2idx, max_len, code='code', age='age'):
+        self.vocab = token2idx
+        self.max_len = max_len
+        self.code = dataframe[code]
+        self.age = dataframe[age]
+        self.age2idx = age2idx
+
+    def __getitem__(self, index):
+        """
+        return: age, code, position, segmentation, mask, label
+        """
+
+        # extract data
+        age = self.age[index][(-self.max_len+1):]
+        code = self.code[index][(-self.max_len+1):]
+
+        # avoid data cut with first element to be 'SEP'
+        if code[0] != 'SEP':
+            code = np.append(np.array(['CLS']), code)
+            age = np.append(np.array(age[0]), age)
+        else:
+            code[0] = 'CLS'
+
+        # mask 0:len(code) to 1, padding to be 0
+        mask = np.ones(self.max_len)
+        mask[len(code):] = 0
+
+        # pad age sequence and code sequence
+        age = seq_padding(age, self.max_len, token2idx=self.age2idx)
+
+        tokens, code, label = random_mask(code, self.vocab)
+
+        # get position code and segment code
+        tokens = seq_padding(tokens, self.max_len)
+        position = position_idx(tokens)
+        segment = index_seg(tokens)
+
+        # pad code and label
+        code = seq_padding(code, self.max_len, symbol=self.vocab['PAD'])
+        label = seq_padding(label, self.max_len, symbol=-1)
+
+        return torch.LongTensor(age), torch.LongTensor(code), torch.LongTensor(position), torch.LongTensor(segment), \
+               torch.LongTensor(mask), torch.LongTensor(label)
+
+    def __len__(self):
+        return len(self.code)
\ No newline at end of file