Function get_datasets

data_preprocessing/prepare_snlr_datasets.py:9–39

(difficulty="easy")

from src.helm.benchmark.scenarios.synthetic_reasoning_natural_scenario import SRNScenario

def get_datasets(difficulty="easy"):
    train = []
    val = []
    test = []

    srn = SRNScenario(difficulty)
    instances = srn.get_instances()

    for instance in instances:
        inputs = instance.input.split("\n")
        fact_index = inputs.index("Fact:")
        rules = inputs[:fact_index]
        fact = inputs[fact_index+1]
        question = inputs[fact_index+2]
        consequents = []
        for ref in instance.references:
            consequents.append(ref.output)
        sample = {
            "rules": rules,
            "fact": fact,
            "question": question,
            "consequents": consequents
        }
        if instance.split == "valid":
            val.append(sample)
        elif instance.split == "test":
            test.append(sample)
        else:
            train.append(sample)
    
    return train, val, test

def save_dataset(dataset, path):
    with open(path, "w") as f: