merge deepspeed

model/supervised_finetuning/custom_datasets/__init__.py

@@ -2,6 +2,8 @@ from datasets import load_dataset
 from sklearn.model_selection import train_test_split
 from torch.utils.data import Dataset, Subset
 
+from .prompt_dialogue import PromptGeneratedDataset
+
 QA_SPECIAL_TOKENS = {"Question": "<question>", "Answer": "<answer>"}
 
 
@@ -14,8 +16,8 @@ class SquadV2Dataset(Dataset):
 
     def __getitem__(self, idx):
         data = self.dataset[idx]
-        # dummy return first answer
-        return "".join([data["title"], ". ", data["context"], " " + data["question"]]), data["answers"]["text"][0]
+        # return first answer form list of possible answers
+        return data["title"] + ". " + data["context"] + " " + data["question"], data["answers"]["text"][0]
 
 
 class WebGPT(Dataset):
@@ -57,12 +59,14 @@ def get_one_dataset(conf, dataset_name):
     dataset_name = dataset_name.lower()
 
     if dataset_name == "squadv2":
-        raise ValueError("SquadV2 is not diverse enough for generation .. ")
         train = SquadV2Dataset(conf.cache_dir, "train")
         eval = SquadV2Dataset(conf.cache_dir, "validation")
     elif dataset_name == "webgpt":
         dataset = WebGPT()
         train, eval = train_val_dataset(dataset, val_split=0.2)
+    elif dataset_name == "prompt_dialogue":
+        dataset = PromptGeneratedDataset()
+        train, eval = train_val_dataset(dataset, val_split=0.2)
     else:
         raise ValueError(f"Unknown dataset {dataset_name}")