| 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"] |
| 55 | data = datasets.load_dataset('gsm8k', name="main", cache_dir=cache_dir) |
| 56 | df = data[type].to_pandas() |
| 57 | sampled_df = df.sample(sample_size, random_state=self.seed)[["question", "answer"]].reset_index(drop=True) |
| 58 | return sampled_df |
| 59 | |
| 60 | def load_physics_question(self, cache_dir="data/bigbench/physics_question.csv", sample_size=None): |
| 61 | df = pd.read_csv(cache_dir) |