Upload dataset.py

updated dataset.py

src/dataset.py

@@ -7,208 +7,18 @@ import random
 from .vocab import Vocab
 import pickle
 import copy
-# from sklearn.preprocessing import OneHotEncoder
+import os
 
-class PretrainerDataset(Dataset):
-    """
-        Class name: PretrainDataset
-        
-    """
-    def __init__(self, dataset_path, vocab, seq_len=30, max_mask=0.15):
-        self.dataset_path = dataset_path
-        self.vocab = vocab # Vocab object
-        
-        # Related to input dataset file
-        self.lines = []
-        self.index_documents = {}
-
-        seq_len_list = []
-        with open(self.dataset_path, "r") as reader:
-            i = 0
-            index = 0
-            self.index_documents[i] = []
-            for line in tqdm.tqdm(reader.readlines()):
-                if line:
-                    line = line.strip()
-                    if not line:
-                        i+=1
-                        self.index_documents[i] = []
-                    else:
-                        self.index_documents[i].append(index)
-                        self.lines.append(line.split("\t"))
-                        len_line = len(line.split("\t"))
-                        seq_len_list.append(len_line)
-                        index+=1
-            reader.close()
-        print("Sequence Stats: len: %s, min: %s, max: %s, average: %s"% (len(seq_len_list),
-              min(seq_len_list), max(seq_len_list), sum(seq_len_list)/len(seq_len_list)))
-        print("Unique Sequences: ", len({tuple(ll) for ll in self.lines}))
-        self.index_documents = {k:v for k,v in self.index_documents.items() if v}
-        print(len(self.index_documents))
-        self.seq_len = seq_len
-        print("Sequence length set at: ", self.seq_len)
-        self.max_mask = max_mask
-        print("% of input tokens selected for masking : ",self.max_mask)
-
-    
-    def __len__(self):
-        return len(self.lines)
-    
-    def __getitem__(self, item):
-        token_a = self.lines[item]
-        # sa_masked = None
-        # sa_masked_label = None
-        # token_b = None
-        # is_same_student = None
-        # sb_masked = None
-        # sb_masked_label = None
-        
-        # if self.select_next_seq:
-        #     is_same_student, token_b = self.get_token_b(item)
-        #     is_same_student = 1 if is_same_student else 0
-        #     token_a1, token_b1 = self.truncate_to_max_seq(token_a, token_b)
-        #     sa_masked, sa_masked_label = self.random_mask_seq(token_a1)
-        #     sb_masked, sb_masked_label = self.random_mask_seq(token_b1)
-        # else:
-        token_a = token_a[:self.seq_len-2]
-        sa_masked, sa_masked_label, sa_masked_pos = self.random_mask_seq(token_a)
-
-        s1 = ([self.vocab.vocab['[CLS]']] + sa_masked + [self.vocab.vocab['[SEP]']])
-        s1_label = ([self.vocab.vocab['[PAD]']] + sa_masked_label + [self.vocab.vocab['[PAD]']])
-        segment_label = [1 for _ in range(len(s1))]
-        masked_pos = ([0] + sa_masked_pos + [0])
-        
-        # if self.select_next_seq:
-        #     s1 = s1 + sb_masked + [self.vocab.vocab['[SEP]']]
-        #     s1_label = s1_label + sb_masked_label + [self.vocab.vocab['[PAD]']]
-        #     segment_label = segment_label + [2 for _ in range(len(sb_masked)+1)]
-        
-        padding = [self.vocab.vocab['[PAD]'] for _ in range(self.seq_len - len(s1))]
-        s1.extend(padding)
-        s1_label.extend(padding)
-        segment_label.extend(padding)
-        masked_pos.extend(padding)
- 
-        output = {'bert_input': s1,
-                 'bert_label': s1_label,
-                 'segment_label': segment_label,
-                 'masked_pos': masked_pos}
-        # print(f"tokenA: {token_a}")
-        # print(f"output: {output}")
-        
-        # if self.select_next_seq:
-        #     output['is_same_student'] = is_same_student
-        
-        # print(item, len(s1), len(s1_label), len(segment_label))
-        # print(f"{item}.")
-        return {key: torch.tensor(value) for key, value in output.items()}
-    
-    def random_mask_seq(self, tokens):
-        """
-        Input: original token seq
-        Output: masked token seq, output label
-        """
-        
-        masked_pos = []
-        output_labels = []
-        output_tokens = copy.deepcopy(tokens)
-        opt_step = False
-        for i, token in enumerate(tokens):
-            if token in ['OptionalTask_1', 'EquationAnswer', 'NumeratorFactor', 'DenominatorFactor', 'OptionalTask_2', 'FirstRow1:1', 'FirstRow1:2', 'FirstRow2:1', 'FirstRow2:2', 'SecondRow', 'ThirdRow']:
-                opt_step = True
-            # if opt_step:
-            #     prob = random.random()
-            #     if prob < self.max_mask:
-            #         output_tokens[i] = random.choice([3,7,8,9,11,12,13,14,15,16,22,23,24,25,26,27,30,31,32])
-            #         masked_pos.append(1)
-            #     else:
-            #         output_tokens[i] = self.vocab.vocab.get(token, self.vocab.vocab['[UNK]'])
-            #         masked_pos.append(0)
-            #     output_labels.append(self.vocab.vocab.get(token, self.vocab.vocab['[UNK]']))
-            #     opt_step = False
-            # else:    
-            prob = random.random()
-            if prob < self.max_mask:
-             # chooses 15% of token positions at random
-                # prob /= 0.15
-                prob = random.random()
-                if prob < 0.8: #[MASK] token 80% of the time
-                    output_tokens[i] = self.vocab.vocab['[MASK]']
-                    masked_pos.append(1)
-                elif prob < 0.9: # a random token 10% of the time 
-                    # print(".......0.8-0.9......")
-                    if opt_step:
-                        output_tokens[i] = random.choice([7,8,9,11,12,13,14,15,16,22,23,24,25,26,27,30,31,32])
-                        opt_step = False
-                    else:
-                        output_tokens[i] = random.randint(1, len(self.vocab.vocab)-1)
-                    masked_pos.append(1)
-                else: # the unchanged i-th token 10% of the time
-                    # print(".......unchanged......")
-                    output_tokens[i] = self.vocab.vocab.get(token, self.vocab.vocab['[UNK]'])
-                    masked_pos.append(0)
-                # True Label
-                output_labels.append(self.vocab.vocab.get(token, self.vocab.vocab['[UNK]']))
-                # masked_pos_label[i] = self.vocab.vocab.get(token, self.vocab.vocab['[UNK]'])
-            else:
-                # i-th token with original value
-                output_tokens[i] = self.vocab.vocab.get(token, self.vocab.vocab['[UNK]'])
-                # Padded label
-                output_labels.append(self.vocab.vocab['[PAD]'])
-                masked_pos.append(0)
-        # label_position = []
-        # label_tokens = []
-        # for k, v in masked_pos_label.items():
-        #     label_position.append(k)
-        #     label_tokens.append(v)
-        return  output_tokens, output_labels, masked_pos
-    
-#     def get_token_b(self, item):
-#         document_id = [k for k,v in self.index_documents.items() if item in v][0]
-#         random_document_id = document_id
-        
-#         if random.random() < 0.5:
-#             document_ids = [k for k in self.index_documents.keys() if k != document_id]
-#             random_document_id = random.choice(document_ids) 
-
-#         same_student = (random_document_id == document_id)
-        
-#         nex_seq_list = self.index_documents.get(random_document_id)
-
-#         if same_student:
-#             if len(nex_seq_list) != 1:
-#                 nex_seq_list = [v for v in nex_seq_list if v !=item]
-
-#         next_seq = random.choice(nex_seq_list)
-#         tokens = self.lines[next_seq]
-#         # print(f"item = {item}, tokens: {tokens}")
-#         # print(f"item={item}, next={next_seq}, same_student = {same_student}, {document_id} == {random_document_id}, b. {tokens}")
-#         return same_student, tokens
-
-#     def truncate_to_max_seq(self, s1, s2):
-#         sa = copy.deepcopy(s1)
-#         sb = copy.deepcopy(s1)
-#         total_allowed_seq = self.seq_len - 3
-        
-#         while((len(sa)+len(sb)) > total_allowed_seq):
-#             if random.random() < 0.5:
-#                 sa.pop()
-#             else:
-#                 sb.pop()
-#         return sa, sb
-            
-                
 class TokenizerDataset(Dataset):
     """
         Class name: TokenizerDataset
         Tokenize the data in the dataset
-        
+        Feat length: 17 
     """
     def __init__(self, dataset_path, label_path, vocab, seq_len=30):
         self.dataset_path = dataset_path
         self.label_path = label_path
         self.vocab = vocab # Vocab object
-        # self.encoder = OneHotEncoder(sparse=False)
         
         # Related to input dataset file
         self.lines = []
@@ -242,39 +52,14 @@ class TokenizerDataset(Dataset):
                         feat2 = [float(i) for i in line.split(",")[-2].split("\t")]
                         feat_vec.extend(feat2[1:])
                         
-                        # # highGRschool_w_prior_w_p_diffskill_wo_fa
-                        # feat_vec = [float(i) for i in line.split(",")[-3].split("\t")]
-                        # feat2 = [-float(i) for i in line.split(",")[-2].split("\t")]
-                        # feat_vec.extend(feat2[1:])
-                        
-#                         # highGRschool_w_prior_w_diffskill_0fa_skill
-#                         feat_vec = [float(i) for i in line.split(",")[-3].split("\t")]
-#                         feat2 = [float(i) for i in line.split(",")[-2].split("\t")]
-#                         fa_feat_vec = [float(i) for i in line.split(",")[-1].split("\t")]
-                        
-#                         diff_skill = [f2 if f1==0 else 0 for f2, f1 in zip(feat2, fa_feat_vec)]
-#                         feat_vec.extend(diff_skill)
-                         
                         if j == 0:
                             print(len(feat_vec))
                             j+=1
-                        
-                        # feat_vec.extend(feat2[1:])
-                        # feat_vec.extend(feat2)
-                        # feat_vec = [float(i) for i in line.split(",")[-2].split("\t")]
-                        # feat_vec = feat_vec[1:]
-                        # feat_vec = [float(line.split(",")[-1])]
-                        # feat_vec = [float(i) for i in line.split(",")[-1].split("\t")]
-                        # feat_vec = [ft-f1 for ft, f1 in zip(feat_vec, fa_feat_vec)]
-
+                    
                         self.feats.append(feat_vec)
                 dataset_info_file.close()
             except Exception as e:
                 print(e)
-            # labeler = np.array([0, 1]) #np.unique(self.labels)
-            # print(f"Labeler {labeler}")
-            # self.encoder.fit(labeler.reshape(-1,1))
-            # self.labels = self.encoder.transform(np.array(self.labels).reshape(-1,1))
 
         self.file = open(self.dataset_path, "r")
         for line in self.file:
@@ -317,95 +102,197 @@ class TokenizerDataset(Dataset):
                  'segment_label': segment_label}
         return {key: torch.tensor(value) for key, value in output.items()}
         
-        
-class TokenizerDatasetForCalibration(Dataset):
+class TokenizerwSkillsDataset(Dataset):
     """
-        Class name: TokenizerDataset
-        Tokenize the data in the dataset
-        
+        Feature length: 17
+
     """
     def __init__(self, dataset_path, label_path, vocab, seq_len=30):
+        print(f"dataset_path: {dataset_path}")
+        print(f"label_path: {label_path}")
+
         self.dataset_path = dataset_path
         self.label_path = label_path
         self.vocab = vocab # Vocab object
-        # self.encoder = OneHotEncoder(sparse=False)
-        
+        self.seq_len = seq_len
+
         # Related to input dataset file
         self.lines = []
         self.labels = []
         self.feats = []
+        selected_lines = []
+
+        print("TokenizerwSkillsDataset...............................")
+
         if self.label_path:
+            # Comment this section if you are not using feat attribute
+            dataset_info_file = open(self.label_path.replace("label", "info"), "r").readlines()
+            print(">>>>>>>>>>>>>>>>>", len(dataset_info_file))
+            j = 0
+            for idex, line in enumerate(dataset_info_file):
+                try:
+                    if line:
+                        line = line.strip()
+                        if not line:
+                            continue
+
+                        feat_vec = [float(i) for i in line.split(",")[-9].split("\t")]
+                        feat2 = [float(i) for i in line.split(",")[-8].split("\t")]
+                        feat_vec.extend(feat2[1:])
+
+                        if j == 0:
+                            print(";;;;", len(feat_vec), feat_vec)
+                            j+=1
+                        self.feats.append(feat_vec)
+                        selected_lines.append(idex)
+                except Exception as e:
+                    print("................>")
+                    print(e)
+                    print("Error at index: ", idex)
+
             self.label_file = open(self.label_path, "r")
-            for line in self.label_file:
+            for idex, line in enumerate(self.label_file):
                 if line:
                     line = line.strip()
                     if not line:
                         continue
-                    self.labels.append(int(line))
+                    if idex in selected_lines:
+                        self.labels.append(int(line))
+                    # self.labels.append(int(line))
             self.label_file.close()
-            
+
+        self.file = open(self.dataset_path, "r")
+        for idex, line in enumerate(self.file):
+            if line:
+                line = line.strip()
+                if line:
+                    if idex in selected_lines:
+                        self.lines.append(line)
+                    # self.lines.append(line)
+        self.file.close()
+        self.len = len(self.lines)
+        print("Sequence length set at ", self.seq_len, len(self.lines), len(self.labels) if self.label_path else 0)
+
+    def __len__(self):
+        return self.len
+
+    def __getitem__(self, item):
+        org_line = self.lines[item].split("\t")
+        dup_line = []
+        opt = False
+        for l in org_line:
+            if l in ["OptionalTask_1", "EquationAnswer", "NumeratorFactor", "DenominatorFactor", "OptionalTask_2", "FirstRow1:1", "FirstRow1:2", "FirstRow2:1", "FirstRow2:2", "SecondRow", "ThirdRow"]:
+                opt = True
+            if opt and 'FinalAnswer-' in l:
+                dup_line.append('[UNK]')
+            else:
+                dup_line.append(l)
+        dup_line = "\t".join(dup_line)
+        # print(dup_line)
+        s1 = self.vocab.to_seq(dup_line, self.seq_len) # This is like tokenizer and adds [CLS] and [SEP].
+        s1_label = self.labels[item] if self.label_path else 0
+        segment_label = [1 for _ in range(len(s1))]
+        s1_feat = self.feats[item] if len(self.feats)>0 else 0
+        padding = [self.vocab.vocab['[PAD]'] for _ in range(self.seq_len - len(s1))]
+        s1.extend(padding), segment_label.extend(padding)
+        # print(s1_feat)
+
+        output = {'input': s1,
+                 'label': s1_label,
+                  'feat': s1_feat,
+                 'segment_label': segment_label}
+        return {key: torch.tensor(value) for key, value in output.items()}
+
+
+class TokenizerwTimeDataset(Dataset):
+    """
+        Feature length: 4
+
+    """
+    def __init__(self, dataset_path, label_path, vocab, seq_len=30):
+        print(f"dataset_path: {dataset_path}")
+        print(f"label_path: {label_path}")
+
+        self.dataset_path = dataset_path
+        self.label_path = label_path
+        self.vocab = vocab # Vocab object
+        self.seq_len = seq_len
+
+        # Related to input dataset file
+        self.lines = []
+        self.labels = []
+        self.feats = []
+        selected_lines = []
+
+        print("TokenizerwTimeDataset...............................")
+        time_df = pickle.load(open("ratio_proportion_change3_2223/sch_largest_100-coded/time_info/full_data_normalized_time.pkl", "rb"))
+        print("time: ?? ", time_df.shape)
+
+        if self.label_path:
             # Comment this section if you are not using feat attribute
-            try:
-                j = 0
-                dataset_info_file = open(self.label_path.replace("label", "info"), "r")
-                for line in dataset_info_file:
+            dataset_info_file = open(self.label_path.replace("label", "info"), "r").readlines()
+            print(">>>>>>>>>>>>>>>>>", len(dataset_info_file))
+            j = 0
+            for idex, line in enumerate(dataset_info_file):
+                try:
                     if line:
                         line = line.strip()
                         if not line:
                             continue
-                      
-                        # # highGRschool_w_prior
-                        # feat_vec = [float(i) for i in line.split(",")[-3].split("\t")]
-                        
-                        # highGRschool_w_prior_w_diffskill_wo_fa
-                        feat_vec = [float(i) for i in line.split(",")[-3].split("\t")]
-                        feat2 = [float(i) for i in line.split(",")[-2].split("\t")]
-                        feat_vec.extend(feat2[1:])
-                        
-#                         # highGRschool_w_prior_w_diffskill_0fa_skill
-#                         feat_vec = [float(i) for i in line.split(",")[-3].split("\t")]
-#                         feat2 = [float(i) for i in line.split(",")[-2].split("\t")]
-#                         fa_feat_vec = [float(i) for i in line.split(",")[-1].split("\t")]
-                        
-#                         diff_skill = [f2 if f1==0 else 0 for f2, f1 in zip(feat2, fa_feat_vec)]
-#                         feat_vec.extend(diff_skill)
-                         
+
+                        feat_vec = []
+
+                        sch = line.split(",")[0]
+                        stu = line.split(",")[2]
+                        progress = line.split(",")[3]
+                        prob_id = line.split(",")[4]
+
+                        total_time = time_df.loc[(sch, stu, progress, prob_id)]['total_time'].item()
+                        faopt_time = time_df.loc[(sch, stu, progress, prob_id)]['faopt_time'].item()
+                        opt_time = time_df.loc[(sch, stu, progress, prob_id)]['opt_time'].item()
+                        nonopt_time = time_df.loc[(sch, stu, progress, prob_id)]['nonopt_time'].item()
+
+                        feat_vec.append(faopt_time)
+                        feat_vec.append(total_time)
+                        feat_vec.append(opt_time)
+                        feat_vec.append(nonopt_time)
+
                         if j == 0:
-                            print(len(feat_vec))
+                            print(";;;;", len(feat_vec), feat_vec)
                             j+=1
-                        
-                        # feat_vec.extend(feat2[1:])
-                        # feat_vec.extend(feat2)
-                        # feat_vec = [float(i) for i in line.split(",")[-2].split("\t")]
-                        # feat_vec = feat_vec[1:]
-                        # feat_vec = [float(line.split(",")[-1])]
-                        # feat_vec = [float(i) for i in line.split(",")[-1].split("\t")]
-                        # feat_vec = [ft-f1 for ft, f1 in zip(feat_vec, fa_feat_vec)]
-
                         self.feats.append(feat_vec)
-                dataset_info_file.close()
-            except Exception as e:
-                print(e)
-            # labeler = np.array([0, 1]) #np.unique(self.labels)
-            # print(f"Labeler {labeler}")
-            # self.encoder.fit(labeler.reshape(-1,1))
-            # self.labels = self.encoder.transform(np.array(self.labels).reshape(-1,1))
+                        selected_lines.append(idex)
+                except Exception as e:
+                    print("................>")
+                    print(e)
+                    print("Error at index: ", idex)
+
+            self.label_file = open(self.label_path, "r")
+            for idex, line in enumerate(self.label_file):
+                if line:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    if idex in selected_lines:
+                        self.labels.append(int(line))
+                    # self.labels.append(int(line))
+            self.label_file.close()
 
         self.file = open(self.dataset_path, "r")
-        for line in self.file:
+        for idex, line in enumerate(self.file):
             if line:
                 line = line.strip()
                 if line:
-                    self.lines.append(line)
-        self.file.close()             
-        
+                    if idex in selected_lines:
+                        self.lines.append(line)
+                    # self.lines.append(line)
+        self.file.close()
         self.len = len(self.lines)
-        self.seq_len = seq_len
         print("Sequence length set at ", self.seq_len, len(self.lines), len(self.labels) if self.label_path else 0)
-        
+
     def __len__(self):
         return self.len
-    
+
     def __getitem__(self, item):
         org_line = self.lines[item].split("\t")
         dup_line = []
@@ -413,7 +300,7 @@ class TokenizerDatasetForCalibration(Dataset):
         for l in org_line:
             if l in ["OptionalTask_1", "EquationAnswer", "NumeratorFactor", "DenominatorFactor", "OptionalTask_2", "FirstRow1:1", "FirstRow1:2", "FirstRow2:1", "FirstRow2:2", "SecondRow", "ThirdRow"]:
                 opt = True
-            if opt and 'FinalAnswer-' in l: 
+            if opt and 'FinalAnswer-' in l:
                 dup_line.append('[UNK]')
             else:
                 dup_line.append(l)
@@ -425,35 +312,128 @@ class TokenizerDatasetForCalibration(Dataset):
         s1_feat = self.feats[item] if len(self.feats)>0 else 0
         padding = [self.vocab.vocab['[PAD]'] for _ in range(self.seq_len - len(s1))]
         s1.extend(padding), segment_label.extend(padding)
-        
+        # print(s1_feat)
+
         output = {'input': s1,
                  'label': s1_label,
                   'feat': s1_feat,
                  'segment_label': segment_label}
-        return ({key: torch.tensor(value) for key, value in output.items()}, s1_label)
-        
-        
-        
-        # if __name__ == "__main__":
-#     # import pickle
-#     # k = pickle.load(open("dataset/CL4999_1920/unique_steps_list.pkl","rb"))
-#     # print(k)
-#     vocab_obj = Vocab("pretraining/vocab.txt")
-#     vocab_obj.load_vocab()
-#     datasetTrain = PretrainerDataset("pretraining/pretrain.txt", vocab_obj)
-    
-#     print(datasetTrain, len(datasetTrain))#, datasetTrain.documents_index)
-#     print(datasetTrain[len(datasetTrain)-1])
-#     for i, d in enumerate(datasetTrain):
-#         print(d.items())
-#         break
-        
-#     fine_tune = TokenizerDataset("finetuning/finetune.txt", "finetuning/finetune_label.txt", vocab_obj)
-#     print(fine_tune)
-#     print(fine_tune[len(fine_tune)-1])
-#     print(fine_tune[random.randint(0, len(fine_tune))])
-#     for i, d in enumerate(fine_tune):
-#         print(d.items())
-#         break
+        return {key: torch.tensor(value) for key, value in output.items()}
         
-    
+class TokenizerwSkillsTimeDataset(Dataset):
+    """
+        Feature length: 17+4 = 21
+
+    """
+    def __init__(self, dataset_path, label_path, vocab, seq_len=30):
+        print(f"dataset_path: {dataset_path}")
+        print(f"label_path: {label_path}")
+
+        self.dataset_path = dataset_path
+        self.label_path = label_path
+        self.vocab = vocab # Vocab object
+        self.seq_len = seq_len
+
+        # Related to input dataset file
+        self.lines = []
+        self.labels = []
+        self.feats = []
+        selected_lines = []
+
+        print("TokenizerwSkillsTimeDataset...............................")
+        time_df = pickle.load(open("ratio_proportion_change3_2223/sch_largest_100-coded/time_info/full_data_normalized_time.pkl", "rb"))
+        print("time: ", time_df.shape)
+
+        if self.label_path:
+            # Comment this section if you are not using feat attribute
+            dataset_info_file = open(self.label_path.replace("label", "info"), "r").readlines()
+            print(">>>>>>>>>>>>>>>>>", len(dataset_info_file))
+            j = 0
+            for idex, line in enumerate(dataset_info_file):
+                try:
+                    if line:
+                        line = line.strip()
+                        if not line:
+                            continue
+
+                        feat_vec = [float(i) for i in line.split(",")[-9].split("\t")]
+                        feat2 = [float(i) for i in line.split(",")[-8].split("\t")]
+                        feat_vec.extend(feat2[1:])
+
+                        sch = line.split(",")[0]
+                        stu = line.split(",")[2]
+                        progress = line.split(",")[3]
+                        prob_id = line.split(",")[4]
+
+                        total_time = time_df.loc[(sch, stu, progress, prob_id)]['total_time'].item()
+                        faopt_time = time_df.loc[(sch, stu, progress, prob_id)]['faopt_time'].item()
+                        opt_time = time_df.loc[(sch, stu, progress, prob_id)]['opt_time'].item()
+                        nonopt_time = time_df.loc[(sch, stu, progress, prob_id)]['nonopt_time'].item()
+
+                        feat_vec.append(faopt_time)
+                        feat_vec.append(total_time)
+                        feat_vec.append(opt_time)
+                        feat_vec.append(nonopt_time)
+
+                        if j == 0:
+                            print(";;;;", len(feat_vec), feat_vec)
+                            j+=1
+                        self.feats.append(feat_vec)
+                        selected_lines.append(idex)
+                except Exception as e:
+                    print("................>")
+                    print(e)
+                    print("Error at index: ", idex)
+
+            self.label_file = open(self.label_path, "r")
+            for idex, line in enumerate(self.label_file):
+                if line:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    if idex in selected_lines:
+                        self.labels.append(int(line))
+                    # self.labels.append(int(line))
+            self.label_file.close()
+
+        self.file = open(self.dataset_path, "r")
+        for idex, line in enumerate(self.file):
+            if line:
+                line = line.strip()
+                if line:
+                    if idex in selected_lines:
+                        self.lines.append(line)
+                    # self.lines.append(line)
+        self.file.close()
+        self.len = len(self.lines)
+        print("Sequence length set at ", self.seq_len, len(self.lines), len(self.labels) if self.label_path else 0)
+
+    def __len__(self):
+        return self.len
+
+    def __getitem__(self, item):
+        org_line = self.lines[item].split("\t")
+        dup_line = []
+        opt = False
+        for l in org_line:
+            if l in ["OptionalTask_1", "EquationAnswer", "NumeratorFactor", "DenominatorFactor", "OptionalTask_2", "FirstRow1:1", "FirstRow1:2", "FirstRow2:1", "FirstRow2:2", "SecondRow", "ThirdRow"]:
+                opt = True
+            if opt and 'FinalAnswer-' in l:
+                dup_line.append('[UNK]')
+            else:
+                dup_line.append(l)
+        dup_line = "\t".join(dup_line)
+        # print(dup_line)
+        s1 = self.vocab.to_seq(dup_line, self.seq_len) # This is like tokenizer and adds [CLS] and [SEP].
+        s1_label = self.labels[item] if self.label_path else 0
+        segment_label = [1 for _ in range(len(s1))]
+        s1_feat = self.feats[item] if len(self.feats)>0 else 0
+        padding = [self.vocab.vocab['[PAD]'] for _ in range(self.seq_len - len(s1))]
+        s1.extend(padding), segment_label.extend(padding)
+        # print(s1_feat)
+
+        output = {'input': s1,
+                 'label': s1_label,
+                  'feat': s1_feat,
+                 'segment_label': segment_label}
+        return {key: torch.tensor(value) for key, value in output.items()}