| 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) | |