bigscience/promptsource

1import time2from jinja2 import meta, TemplateError3import pytest4import promptsource.templates5from promptsource.utils import get_dataset_builder6from uuid import UUID7 8# Sets up Jinja environment9env = promptsource.templates.env10 11# Loads templates and iterates over each data (sub)set12template_collection = promptsource.templates.TemplateCollection()13 14 15def test_uuids():16    """17    Checks that all UUIDs across promptsource are unique. (Although collisions18    are unlikely, copying and pasting YAML files could lead to duplicates.19    """20    all_uuids = {}21 22    # Iterates over all datasets23    for dataset_name, subset_name in template_collection.keys:24 25        # Iterates over each template for current data (sub)set26        dataset_templates = template_collection.get_dataset(dataset_name, subset_name)27        for template_name in dataset_templates.all_template_names:28            template = dataset_templates[template_name]29 30            uuid = template.get_id()31 32            if uuid in all_uuids:33                raise ValueError(f"Template {template_name} for dataset {dataset_name}/{subset_name} "34                                 f"has duplicate uuid {template.get_id()} as "35                                 f"{all_uuids[uuid][0]}/{all_uuids[uuid][1]}.")36 37            all_uuids[uuid] = (dataset_name, subset_name)38 39 40@pytest.mark.parametrize("dataset", template_collection.keys)41def test_dataset(dataset):42    """43    Validates all the templates in the repository with simple syntactic checks:44    0. Are all templates parsable YAML?45    1. Do all templates parse in Jinja and are all referenced variables in the dataset schema?46    2. Does the template contain a prompt/output separator "|||" ?47    3. Are all names and templates within a data (sub)set unique?48    4. Is the YAML dictionary properly formatted?49    5. Is the UUID valid?50 51    :param dataset: (dataset_name, subset_name) pair to test52 53    """54    dataset_name, subset_name = dataset55 56    # Loads dataset information57    tries = 058    max_tries = 359    while True:60        try:61            builder_instance = get_dataset_builder(dataset_name, subset_name)62            break63        except ConnectionError as e:64            if tries < max_tries:65                time.sleep(2)66                tries += 167            else:68                raise e69 70    has_features = builder_instance.info.features is not None71    if has_features:72        features = builder_instance.info.features.keys()73        features = set([feature.replace("-", "_") for feature in features])74 75    # Initializes sets for checking uniqueness among templates76    template_name_set = set()77    template_jinja_set = set()78 79    # Iterates over each template for current data (sub)set80    dataset_templates = template_collection.get_dataset(dataset_name, subset_name)81    any_original = False82    for template_name in dataset_templates.all_template_names:83        template = dataset_templates[template_name]84        any_original = any_original or template.metadata.original_task85        # Check 1: Jinja and all features valid?86        try:87            parse = env.parse(template.jinja)88        except TemplateError as e:89            raise ValueError(f"Template for dataset {dataset_name}/{subset_name} "90                             f"with uuid {template.get_id()} failed to parse.") from e91 92        variables = meta.find_undeclared_variables(parse)93        for variable in variables:94            if has_features and variable not in features and variable != "answer_choices":95                raise ValueError(f"Template for dataset {dataset_name}/{subset_name} "96                                 f"with uuid {template.get_id()} has unrecognized variable {variable}.")97 98        # Check 2: Prompt/output separator present?99        if "|||" not in template.jinja:100            raise ValueError(f"Template for dataset {dataset_name}/{subset_name} "101                             f"with uuid {template.get_id()} has no prompt/output separator.")102 103        # Check 3: Unique names and templates?104        if template.get_name() in template_name_set:105            raise ValueError(f"Template for dataset {dataset_name}/{subset_name} "106                             f"with uuid {template.get_id()} has duplicate name.")107 108        if template.jinja in template_jinja_set:109            raise ValueError(f"Template for dataset {dataset_name}/{subset_name} "110                             f"with uuid {template.get_id()} has duplicate definition.")111 112        template_name_set.add(template.get_name())113        template_jinja_set.add(template.jinja)114 115        # Check 4: Is the YAML dictionary properly formatted?116        try:117            if dataset_templates.templates[template.get_id()] != template:118                raise ValueError(f"Template for dataset {dataset_name}/{subset_name} "119                                 f"with uuid {template.get_id()} has wrong YAML key.")120        except KeyError as e:121            raise ValueError(f"Template for dataset {dataset_name}/{subset_name} "122                             f"with uuid {template.get_id()} has wrong YAML key.") from e123 124        # Check 5: Is the UUID valid?125        UUID(template.get_id())126 127    # Turned off for now until we fix.128    #assert any_original, "There must be at least one original task template for each dataset"129