| 44 | return sampled_df |
| 45 | |
| 46 | def load_trivia_qa(self, cache_dir="data/trivia_qa", sample_size=100, type="test"): |
| 47 | assert type in ["train", "validation", "test"] |
| 48 | data = datasets.load_dataset('trivia_qa', 'rc.nocontext', cache_dir=cache_dir) |
| 49 | df = data[type].to_pandas() |
| 50 | sampled_df = df.sample(sample_size, random_state=self.seed)[["question", "answer"]].reset_index(drop=True) |
| 51 | return sampled_df |
| 52 | |
| 53 | def load_gsm8k(self, cache_dir="data/gsm8k", sample_size=100, type="test"): |
| 54 | assert type in ["train", "validation", "test"] |