(dataset_name: str, n_prompts: int)
| 22 | |
| 23 | |
| 24 | def get_prompts_text(dataset_name: str, n_prompts: int) -> Optional[list[str]]: |
| 25 | ret = [] |
| 26 | if dataset_name.lower() == "mmlu": |
| 27 | logger.info("Loading MMLU dataset...") |
| 28 | ret = datasets.load_dataset("cais/mmlu", "all")["test"]["question"] # type: ignore |
| 29 | else: |
| 30 | return None |
| 31 | if n_prompts >= 0: |
| 32 | ret = ret[:n_prompts] |
| 33 | return ret |
| 34 | |
| 35 | |
| 36 | def get_prompt_lengths_rng(n_prompts: int, prompt_length_min: int, prompt_length_max: int, seed_offset: int) -> list[int]: |