| 221 | |
| 222 | |
| 223 | class PromptDataset(Dataset): |
| 224 | |
| 225 | def __init__(self, prompt_dataset, answer_dataset) -> None: |
| 226 | super().__init__() |
| 227 | self.prompt_dataset = prompt_dataset |
| 228 | self.answer_dataset = answer_dataset |
| 229 | assert len(self.prompt_dataset) == len(self.answer_dataset) |
| 230 | |
| 231 | def __len__(self): |
| 232 | return len(self.prompt_dataset) |
| 233 | |
| 234 | def __getitem__(self, idx): |
| 235 | return { |
| 236 | "prompt": self.prompt_dataset[idx], |
| 237 | "answer": self.answer_dataset[idx] |
| 238 | } |
| 239 | |
| 240 | |
| 241 | def get_prompt_dataset(current_dataset, raw_dataset, add_sys_prefix=False, sample_ratio=None): |
no outgoing calls
no test coverage detected