File size: 1,714 Bytes
9213f5f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
from datasets import load_dataset
import pandas as pd
from tqdm.notebook import tqdm
from sklearn.utils import shuffle
train_dataset = load_dataset("squad", split="train")
valid_dataset = load_dataset("squad", split="validation")
df_train = pd.DataFrame(columns=["context", "answer", "question"])
df_validation = pd.DataFrame(columns=["context", "answer", "question"])
count_long = 0
count_short = 0
for index, val in enumerate(tqdm(train_dataset)):
print(index)
passage = val["context"]
question = val["question"]
answer = val["answers"]["text"][0]
no_of_words = len(answer.split())
if no_of_words >= 7:
count_long = count_long + 1
continue
else:
df_train.loc[count_short] = [passage] + [answer] + [question]
count_short = count_short + 1
print("count_long train dataset: ", count_long)
print("count_short train dataset: ", count_short)
count_long = 0
count_short = 0
for index, val in enumerate(tqdm(valid_dataset)):
print(index)
passage = val["context"]
question = val["question"]
answer = val["answers"]["text"][0]
no_of_words = len(answer.split())
if no_of_words >= 7:
count_long = count_long + 1
continue
else:
df_validation.loc[count_short] = [passage] + [answer] + [question]
count_short = count_short + 1
print("count_long validation dataset: ", count_long)
print("count_short validation dataset: ", count_short)
df_train = shuffle(df_train)
df_validation = shuffle(df_validation)
train_save_path = "squad_t5_train.csv"
validation_save_path = "squad_t5_validaton.csv"
df_train.to_csv(train_save_path, index=False)
df_validation.to_csv(validation_save_path, index=False)
|