| 77 | |
| 78 | |
| 79 | class TextPromptDataset(Dataset): |
| 80 | def __init__(self, dataset, split="train"): |
| 81 | self.file_path = os.path.join(dataset, f"{split}.txt") |
| 82 | with open(self.file_path, "r") as f: |
| 83 | self.prompts = [line.strip() for line in f.readlines()] |
| 84 | |
| 85 | def __len__(self): |
| 86 | return len(self.prompts) |
| 87 | |
| 88 | def __getitem__(self, idx): |
| 89 | return {"prompt": self.prompts[idx], "metadata": {}} |
| 90 | |
| 91 | @staticmethod |
| 92 | def collate_fn(examples): |
| 93 | prompts = [example["prompt"] for example in examples] |
| 94 | metadatas = [example["metadata"] for example in examples] |
| 95 | return prompts, metadatas |
| 96 | |
| 97 | |
| 98 | class GenevalPromptDataset(Dataset): |