import torch
from torch.utils.data import Dataset
import pandas as pd
import numpy as np
import tqdm
import random
from .vocab import Vocab
import pickle
import copy
# from sklearn.preprocessing import OneHotEncoder
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
"""
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 = []
self.labels = []
self.feats = []
if self.label_path:
self.label_file = open(self.label_path, "r")
for line in self.label_file:
if line:
line = line.strip()
if not line:
continue
self.labels.append(int(line))
self.label_file.close()
# 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:
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_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:
if line:
line = line.strip()
if 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 = []
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)
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 TokenizerDatasetForCalibration(Dataset):
"""
Class name: TokenizerDataset
Tokenize the data in the dataset
"""
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 = []
self.labels = []
self.feats = []
if self.label_path:
self.label_file = open(self.label_path, "r")
for line in self.label_file:
if line:
line = line.strip()
if not line:
continue
self.labels.append(int(line))
self.label_file.close()
# 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:
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)
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:
if line:
line = line.strip()
if 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 = []
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)
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