add: dataset gen and training scripts

utils/dataset_gen.py

@@ -0,0 +1,61 @@
+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)

utils/requirements.txt

@@ -0,0 +1,7 @@
+datasets==2.14.5
+pandas==2.1.1
+pytorch_lightning==2.1.0
+scikit_learn==1.3.1
+torch==2.1.0
+tqdm==4.66.1
+transformers==4.34.0

utils/t5_train_model.py

@@ -0,0 +1,220 @@
+import argparse
+
+import pandas as pd
+from torch.utils.data import Dataset, DataLoader
+
+from transformers import AdamW, T5ForConditionalGeneration, T5Tokenizer
+
+from tqdm.notebook import tqdm
+import copy
+import pytorch_lightning as pl
+
+
+class QuestionGenerationDataset(Dataset):
+    def __init__(self, tokenizer, filepath, max_len_inp=512, max_len_out=96):
+        self.path = filepath
+
+        self.passage_column = "context"
+        self.answer = "answer"
+        self.question = "question"
+
+        # self.data = pd.read_csv(self.path)
+        self.data = pd.read_csv(self.path, nrows=1000)
+
+        self.max_len_input = max_len_inp
+        self.max_len_output = max_len_out
+        self.tokenizer = tokenizer
+        self.inputs = []
+        self.targets = []
+        self.skippedcount = 0
+        self._build()
+
+    def __len__(self):
+        return len(self.inputs)
+
+    def __getitem__(self, index):
+        source_ids = self.inputs[index]["input_ids"].squeeze()
+        target_ids = self.targets[index]["input_ids"].squeeze()
+
+        src_mask = self.inputs[index][
+            "attention_mask"
+        ].squeeze()  # might need to squeeze
+        target_mask = self.targets[index][
+            "attention_mask"
+        ].squeeze()  # might need to squeeze
+
+        labels = copy.deepcopy(target_ids)
+        labels[labels == 0] = -100
+
+        return {
+            "source_ids": source_ids,
+            "source_mask": src_mask,
+            "target_ids": target_ids,
+            "target_mask": target_mask,
+            "labels": labels,
+        }
+
+    def _build(self):
+        for idx in tqdm(range(len(self.data))):
+            passage, answer, target = (
+                self.data.loc[idx, self.passage_column],
+                self.data.loc[idx, self.answer],
+                self.data.loc[idx, self.question],
+            )
+
+            input_ = "context: %s  answer: %s </s>" % (passage, answer)
+            target = "question: %s </s>" % (str(target))
+
+            # get encoding length of input. If it is greater than self.max_len skip it
+            test_input_encoding = self.tokenizer.encode_plus(
+                input_, truncation=False, return_tensors="pt"
+            )
+
+            length_of_input_encoding = len(test_input_encoding["input_ids"][0])
+
+            if length_of_input_encoding > self.max_len_input:
+                self.skippedcount = self.skippedcount + 1
+                continue
+
+            # tokenize inputs
+            tokenized_inputs = self.tokenizer.batch_encode_plus(
+                [input_],
+                max_length=self.max_len_input,
+                pad_to_max_length=True,
+                return_tensors="pt",
+            )
+            # tokenize targets
+            tokenized_targets = self.tokenizer.batch_encode_plus(
+                [target],
+                max_length=self.max_len_output,
+                pad_to_max_length=True,
+                return_tensors="pt",
+            )
+
+            self.inputs.append(tokenized_inputs)
+            self.targets.append(tokenized_targets)
+
+
+class T5FineTuner(pl.LightningModule):
+    def __init__(self, hparams, t5model, t5tokenizer):
+        super(T5FineTuner, self).__init__()
+        self.save_hyperparameters(hparams)
+        # self.hparams = hparams
+        self.model = t5model
+        self.tokenizer = t5tokenizer
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        decoder_input_ids=None,
+        decoder_attention_mask=None,
+        lm_labels=None,
+    ):
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            decoder_attention_mask=decoder_attention_mask,
+            labels=lm_labels,
+        )
+
+        return outputs
+
+    def training_step(self, batch, batch_idx):
+        outputs = self.forward(
+            input_ids=batch["source_ids"],
+            attention_mask=batch["source_mask"],
+            decoder_input_ids=batch["target_ids"],
+            decoder_attention_mask=batch["target_mask"],
+            lm_labels=batch["labels"],
+        )
+
+        loss = outputs[0]
+        self.log("train_loss", loss)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        outputs = self.forward(
+            input_ids=batch["source_ids"],
+            attention_mask=batch["source_mask"],
+            decoder_input_ids=batch["target_ids"],
+            decoder_attention_mask=batch["target_mask"],
+            lm_labels=batch["labels"],
+        )
+
+        loss = outputs[0]
+        self.log("val_loss", loss)
+        return loss
+
+    def train_dataloader(self):
+        return DataLoader(
+            train_dataset, batch_size=self.hparams.batch_size, num_workers=4
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            validation_dataset, batch_size=self.hparams.batch_size, num_workers=4
+        )
+
+    def configure_optimizers(self):
+        optimizer = AdamW(self.parameters(), lr=3e-4, eps=1e-8)
+        return optimizer
+
+
+if __name__ == "__main__":
+    pl.seed_everything(42)
+    train_file_path = "question_generator/dataset/squad_t5_train.csv"
+    validation_file_path = "question_generator/dataset/squad_t5_validaton.csv"
+
+    t5_tokenizer = T5Tokenizer.from_pretrained("t5-base")
+    t5_model = T5ForConditionalGeneration.from_pretrained("t5-base")
+
+    sample_encoding = t5_tokenizer.encode_plus(
+        "My name is Pipe San Martin",
+        max_length=64,
+        pad_to_max_length=True,
+        truncation=True,
+        return_tensors="pt",
+    )
+
+    print(sample_encoding.keys())
+    print(sample_encoding["input_ids"].shape)
+    print(sample_encoding["input_ids"].squeeze().shape)
+    print(sample_encoding["input_ids"])
+    tokenized_output = t5_tokenizer.convert_ids_to_tokens(
+        sample_encoding["input_ids"].squeeze()
+    )
+    print(f"Tokenized output: {tokenized_output}")
+    decoded_output = t5_tokenizer.decode(
+        sample_encoding["input_ids"].squeeze(),
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True,
+    )
+    print(f"Decoded output: {decoded_output}")
+    train_dataset = QuestionGenerationDataset(t5_tokenizer, train_file_path)
+
+    train_sample = train_dataset[50]
+    decoded_train_input = t5_tokenizer.decode(train_sample["source_ids"])
+    decoded_train_output = t5_tokenizer.decode(train_sample["target_ids"])
+
+    print(decoded_train_input)
+    print(decoded_train_output)
+
+    validation_dataset = QuestionGenerationDataset(t5_tokenizer, validation_file_path)
+    args_dict = dict(
+        batch_size=4,
+    )
+
+    args = argparse.Namespace(**args_dict)
+
+    model = T5FineTuner(args, t5_model, t5_tokenizer)
+
+    trainer = pl.Trainer(max_epochs=1)
+
+    trainer.fit(model)
+
+    #print("Saving model")
+    #save_path_model = "question_generator/model/"
+    #save_path_tokenizer = "question_generator/tokenizer/"
+    #model.model.save_pretrained(save_path_model)
+    #t5_tokenizer.save_pretrained(save_path_tokenizer)