from datasets import load_dataset


def load_problems(dataset, starting_index=None, end_index=None):
    """
    Load a specified number of problems from a dataset starting from a given index.
    
    Args:
        dataset (str): The name of the dataset to load.
        starting_index (int): The index to start loading problems from.
        num_problems (int): The number of problems to load.
        
    Returns:
        list: A list of problems loaded from the dataset.
    """
    # Load the dataset
    if dataset == "math500":
        if starting_index is None:
            starting_index = 30
        if end_index is None:
            end_index = 500
        problems = load_dataset("HuggingFaceH4/MATH-500")["test"]["problem"][starting_index:end_index]
    
    elif dataset == "gsm8k":
        if starting_index is None:
            starting_index = 30
        if end_index is None:
            end_index = 330
        problems = load_dataset("openai/gsm8k", "main")["train"]["question"][starting_index:end_index]

    else:
        raise ValueError("Unsupported dataset. Please use 'math500' or 'gsm8k'.")
    
    return problems