utils/json_loader.py

import json


class JsonDataLoader:
    def __init__(self, filepath):
        self.filepath = filepath

    def load_data(self, category=None):
        with open(self.filepath, "r") as f:
            data = json.load(f)

        inputs = []
        outputs = []

        if category == "easy" or category == "medium":
            inputs += [
                example["inputs"] for example in data if example["category"] == "easy"
            ]
            outputs += [
                example["outputs"] for example in data if example["category"] == "easy"
            ]

        if category == "medium":
            inputs += [
                example["inputs"] for example in data if example["category"] == "medium"
            ]
            outputs += [
                example["outputs"]
                for example in data
                if example["category"] == "medium"
            ]

        if category == None or category == "hard":
            inputs += [example["inputs"] for example in data]
            outputs += [example["outputs"] for example in data]

        return inputs, outputs