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

model/supervised_finetuning/custom_datasets/prompt_dialogue.py

@@ -0,0 +1,66 @@
+import os
+from urllib.request import urlopen
+
+from torch.utils.data import Dataset
+
+
+class PromptGeneratedDataset(Dataset):
+    """Generates from flan 11B
+    User: What are the best methods for preventing a slave trade?
+
+    Rosey: The best methods ....
+    <|endoftext|>
+
+    we are ignoring results with multiple lines for now
+    """
+
+    url = "https://github.com/Rallio67/language-model-agents/raw/main/chat_dialogue_v2_c.txt"
+
+    def __init__(self) -> None:
+        super().__init__()
+        os.makedirs("datasets", exist_ok=True)
+        chat_dialogue = os.path.join("datasets", "chat_dialogue_v2_c.txt")
+        if not os.path.exists(chat_dialogue):
+            with urlopen(self.url) as file:
+                content = file.read().decode()
+            with open(chat_dialogue, "w") as fout:
+                fout.write(content)
+
+        question = ""
+        answer = ""
+        self.pairs = []
+        with open(chat_dialogue, "r") as f:
+            corpus = f.read().split("<|endoftext|>")
+            for dialogue in corpus:
+                dialogue = dialogue.strip()
+                if "Rosey:" in dialogue:
+                    user, bot = dialogue.split("Rosey:", maxsplit=1)
+                    question = user.split(":", maxsplit=1)[1].strip()
+                    answer = bot.strip()
+                    if len(answer) and len(question):
+                        self.pairs.append((question, answer))
+
+        if len(question) > 0 and len(answer) > 0:
+            self.pairs.append((question, answer))
+
+    def __len__(self):
+        return len(self.pairs)
+
+    def __getitem__(self, index):
+        question, answer = self.pairs[index]
+        return question, answer
+
+
+if __name__ == "__main__":
+    from torch.utils.data import DataLoader
+    from transformers import AutoTokenizer
+
+    from .dialogue_collator import DialogueDataCollator
+
+    tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-2B-multi")
+    tokenizer.add_special_tokens({"pad_token": "<|endoftext|>", "sep_token": "<|endoftext|>"})
+    dataset = PromptGeneratedDataset()
+    collate_fn = DialogueDataCollator(tokenizer, padding=True, max_length=128)
+    dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=5)
+    for batch in dataloader:
+        print(batch["input_ids"].shape)