Python _jsonnet.evaluate_snippet() Examples

def render_jsonnet(self, manifeststr, tla_codes=None):
        # @TODO(ant31): workaround until jsonnet compile on windows
        import _jsonnet
        try:
            json_str = _jsonnet.evaluate_snippet(  # pylint: disable=no-member
                "snippet", manifeststr, import_callback=self.import_callback,
                native_callbacks=filters.jsonnet_callbacks(), tla_codes=tla_codes)

        except RuntimeError as e:
            print("tla_codes: %s" % (str(tla_codes)))
            print("\n".join([
                "%s %s" % (i, line) for i, line in enumerate([
                    l for l in manifeststr.split("\n") if re.match(r"^ *#", l) is None])]))
            raise e
        return json.loads(json_str) 

def jsonnet_loads(jsonnet_str, ext_vars=None):
    """
    Parses jsonnet string into json
    :param jsonnet_str: Jsonnet function
    :param ext_vars: External vars that can be passed as {'SOME_PARAM': 'AI2'} and used in the jsonnet as {name: std.extVar("SOME_PARAM")}
    :return:
    """
    json_parse = json.loads(_jsonnet.evaluate_snippet("snippet", jsonnet_str, ext_vars=ext_vars))

    return json_parse 

def evaluate_snippet(_filename: str, expr: str, **_kwargs) -> str:
        logger.warning(
            "error loading _jsonnet (this is expected on Windows), treating snippet as plain json"
        )
        return expr 

def parse_overrides(serialized_overrides: str) -> Dict[str, Any]:
    if serialized_overrides:
        ext_vars = _environment_variables()

        return unflatten(json.loads(evaluate_snippet("", serialized_overrides, ext_vars=ext_vars)))
    else:
        return {} 

def test_regex_matches_are_initialized_correctly(self):
        class Net(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear_1_with_funky_name = torch.nn.Linear(5, 10)
                self.linear_2 = torch.nn.Linear(10, 5)
                self.conv = torch.nn.Conv1d(5, 5, 5)

            def forward(self, inputs):
                pass

        # Make sure we handle regexes properly
        json_params = """{"initializer": {"regexes": [
        ["conv", {"type": "constant", "val": 5}],
        ["funky_na.*bi", {"type": "constant", "val": 7}]
        ]}}
        """
        params = Params(json.loads(_jsonnet.evaluate_snippet("", json_params)))
        initializers = InitializerApplicator.from_params(params=params["initializer"])
        model = Net()
        initializers(model)

        for parameter in model.conv.parameters():
            assert torch.equal(parameter.data, torch.ones(parameter.size()) * 5)

        parameter = model.linear_1_with_funky_name.bias
        assert torch.equal(parameter.data, torch.ones(parameter.size()) * 7) 

def test_regex_match_prevention_prevents_and_overrides(self):
        class Net(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear_1 = torch.nn.Linear(5, 10)
                self.linear_2 = torch.nn.Linear(10, 5)
                # typical actual usage: modules loaded from allenlp.model.load(..)
                self.linear_3_transfer = torch.nn.Linear(5, 10)
                self.linear_4_transfer = torch.nn.Linear(10, 5)
                self.pretrained_conv = torch.nn.Conv1d(5, 5, 5)

            def forward(self, inputs):
                pass

        json_params = """{"initializer": {
        "regexes": [
            [".*linear.*", {"type": "constant", "val": 10}],
            [".*conv.*", {"type": "constant", "val": 10}]
            ],
        "prevent_regexes": [".*_transfer.*", ".*pretrained.*"]
        }}
        """
        params = Params(json.loads(_jsonnet.evaluate_snippet("", json_params)))
        initializers = InitializerApplicator.from_params(params=params["initializer"])
        model = Net()
        initializers(model)

        for module in [model.linear_1, model.linear_2]:
            for parameter in module.parameters():
                assert torch.equal(parameter.data, torch.ones(parameter.size()) * 10)

        transfered_modules = [
            model.linear_3_transfer,
            model.linear_4_transfer,
            model.pretrained_conv,
        ]

        for module in transfered_modules:
            for parameter in module.parameters():
                assert not torch.equal(parameter.data, torch.ones(parameter.size()) * 10) 

def evaluate_snippet(self, *args, **kwargs):
    kwargs['import_callback'] = self._import_callback
    kwargs['native_callbacks'] = self._native_callbacks
    return json.loads(_jsonnet.evaluate_snippet(*args, **kwargs)) 

def test_regex_matches_are_initialized_correctly(self):
        class Net(torch.nn.Module):
            def __init__(self):
                super(Net, self).__init__()
                self.linear_1_with_funky_name = torch.nn.Linear(5, 10)
                self.linear_2 = torch.nn.Linear(10, 5)
                self.conv = torch.nn.Conv1d(5, 5, 5)

            def forward(self, inputs):  # pylint: disable=arguments-differ
                pass

        # Make sure we handle regexes properly
        json_params = u"""{"initializer": [
        ["conv", {"type": "constant", "val": 5}],
        ["funky_na.*bi", {"type": "constant", "val": 7}]
        ]}
        """
        params = Params(json.loads(_jsonnet.evaluate_snippet(u"", json_params)))
        initializers = InitializerApplicator.from_params(params[u'initializer'])
        model = Net()
        initializers(model)

        for parameter in model.conv.parameters():
            assert torch.equal(parameter.data, torch.ones(parameter.size()) * 5)

        parameter = model.linear_1_with_funky_name.bias
        assert torch.equal(parameter.data, torch.ones(parameter.size()) * 7) 

def test_regex_match_prevention_prevents_and_overrides(self):

        class Net(torch.nn.Module):
            def __init__(self):
                super(Net, self).__init__()
                self.linear_1 = torch.nn.Linear(5, 10)
                self.linear_2 = torch.nn.Linear(10, 5)
                # typical actual usage: modules loaded from allenlp.model.load(..)
                self.linear_3_transfer = torch.nn.Linear(5, 10)
                self.linear_4_transfer = torch.nn.Linear(10, 5)
                self.pretrained_conv = torch.nn.Conv1d(5, 5, 5)
            def forward(self, inputs):  # pylint: disable=arguments-differ
                pass

        json_params = u"""{"initializer": [
        [".*linear.*", {"type": "constant", "val": 10}],
        [".*conv.*", {"type": "constant", "val": 10}],
        [".*_transfer.*", "prevent"],
        [".*pretrained.*",{"type": "prevent"}]
        ]}
        """
        params = Params(json.loads(_jsonnet.evaluate_snippet(u"", json_params)))
        initializers = InitializerApplicator.from_params(params[u'initializer'])
        model = Net()
        initializers(model)

        for module in [model.linear_1, model.linear_2]:
            for parameter in module.parameters():
                assert torch.equal(parameter.data, torch.ones(parameter.size())*10)

        transfered_modules = [model.linear_3_transfer,
                              model.linear_4_transfer,
                              model.pretrained_conv]

        for module in transfered_modules:
            for parameter in module.parameters():
                assert not torch.equal(parameter.data, torch.ones(parameter.size())*10) 

def evaluate_snippet(_filename     , expr     , **_kwargs)       :
        logger.warning("_jsonnet not loaded, treating snippet as json")
        return expr 

def parse_overrides(serialized_overrides     )                  :
    if serialized_overrides:
        ext_vars = dict(os.environ)
        return unflatten(json.loads(evaluate_snippet(u"", serialized_overrides, ext_vars=ext_vars)))
    else:
        return {} 

def evaluate_snippet(_filename: str, expr: str, **_kwargs) -> str:
        logger.warning(f"_jsonnet not loaded, treating snippet as json")
        return expr 

def parse_overrides(serialized_overrides: str) -> Dict[str, Any]:
    if serialized_overrides:
        ext_vars = _environment_variables()

        return unflatten(json.loads(evaluate_snippet("", serialized_overrides, ext_vars=ext_vars)))
    else:
        return {} 

def render_jsonnet(self, manifeststr, tla_codes=None):
        try:
            json_str = _jsonnet.evaluate_snippet(
                "snippet", manifeststr, import_callback=self.import_callback,
                native_callbacks=filters.jsonnet_callbacks(), tla_codes=tla_codes)

        except RuntimeError as e:
            print("tla_codes: %s" % (str(tla_codes)))
            print("\n".join([
                "%s %s" % (i, line) for i, line in enumerate(
                    [l for l in manifeststr.split("\n") if re.match(r"^ *#", l) is None])
            ]))
            raise e
        return json.loads(json_str) 

def get_run_func(
        self,
        args: argparse.Namespace,
    ):
        if args is None:
            raise ValueError("No run arguments found for AllenNLP runner.")

        with open(args.base_config, "r") as parameter_f:
            parameter_file_snippet = parameter_f.read()

        def train_func(config, reporter):
            logger.debug(f"CUDA_VISIBLE_DEVICES: {os.environ['CUDA_VISIBLE_DEVICES']}")
            
            for package_name in getattr(args, "include_package", ()):
                import_module_and_submodules(package_name)

            search_space = HyperparameterSearch(**config)
            sample = search_space.sample()
            for k, v in sample.items():
                config[k] = str(v)
            
            params_dict = json.loads(
                _jsonnet.evaluate_snippet(
                    "config", parameter_file_snippet, tla_codes={}, ext_vars=config
                )
            )
            if args.num_gpus == 0:
                logger.warning(f"No GPU specified, using CPU.")
                params_dict["trainer"]["cuda_device"] = -1

            if args.cpus_per_trial > 0:
                torch.set_num_threads(args.cpus_per_trial)

            params = Params(params_dict)

            logger.debug(f"AllenNLP Configuration: {params.as_dict()}")

            train_model(params=params, serialization_dir="trial")

            reporter(done=True)
            
        return train_func