"""Definition of the narchi json schemas."""
import json
import enum
from copy import deepcopy
from jsonschema import Draft7Validator as jsonvalidator
id_pattern = '[A-Za-z_][0-9A-Za-z_]*'
id_separator = 'ยท'
propagated_id_pattern = '[A-Za-z_][0-9A-Za-z_'+id_separator+']*'
variable_pattern = '<<variable:([-+/*0-9A-Za-z_]+)>>'
auto_tag = '<<auto>>'
id_type = {
'type': 'string',
'pattern': '^'+id_pattern+'$',
}
description_type = {
'type': 'string',
'minLength': 8,
'pattern': '^[^<>]+$',
}
dims_type = {
'type': 'array',
'minItems': 1,
'items': {
'oneOf': [
{
'type': 'integer',
'minimum': 1,
},
{
'type': 'string',
'pattern': '^('+variable_pattern+'|'+auto_tag+')$',
},
],
},
}
dims_in_type = deepcopy(dims_type)
dims_in_type['items']['oneOf'].append({'type': 'null'})
shape_type = {
'type': 'object',
'properties': {
'in': {'$ref': '#/definitions/dims_in'},
'out': {'$ref': '#/definitions/dims'},
},
'required': ['in', 'out'],
'additionalProperties': False,
}
graph_type = {
'type': 'array',
'minItems': 1,
'items': {
'type': 'string',
'pattern': '^'+id_pattern+'( +-> +'+id_pattern+')+$',
},
}
path_type = {
'type': 'string',
'pattern': '.+\\.jsonnet',
}
reshape_dims_type = deepcopy(dims_type)
reshape_dims_type['minItems'] = 2
reshape_index_type = {
'type': 'integer',
'minimum': 0,
}
reshape_flatten_type = {
'type': 'array',
'minItems': 2,
'items': {'$ref': '#/definitions/reshape_index'},
}
reshape_unflatten_type = {
'type': 'object',
'minProperties': 1,
'maxProperties': 1,
'patternProperties': {
'^[0-9]+$': {'$ref': '#/definitions/reshape_dims'},
},
'additionalProperties': False,
}
reshape_type = {
'oneOf': [
{
'const': 'flatten',
},
{
'type': 'array',
'minItems': 1,
'items': {
'oneOf': [
{'$ref': '#/definitions/reshape_index'},
{'$ref': '#/definitions/reshape_flatten'},
{'$ref': '#/definitions/reshape_unflatten'},
],
},
},
],
}
reshape_definitions = {
'reshape': reshape_type,
'reshape_index': reshape_index_type,
'reshape_dims': reshape_dims_type,
'reshape_flatten': reshape_flatten_type,
'reshape_unflatten': reshape_unflatten_type,
}
reshape_schema = {
'$ref': '#/definitions/reshape',
'definitions': reshape_definitions,
}
block_type = {
'type': 'object',
'properties': {
'_class': {'$ref': '#/definitions/id'},
'_name': {'$ref': '#/definitions/id'},
'_id': {'$ref': '#/definitions/id'},
'_id_share': {'$ref': '#/definitions/id'},
'_description': {'$ref': '#/definitions/description'},
'_shape': {'$ref': '#/definitions/shape'},
'_path': {'$ref': '#/definitions/path'},
'_ext_vars': {'type': 'object'},
'blocks': {'$ref': '#/definitions/blocks'},
'input': {'$ref': '#/definitions/id'},
'output': {'$ref': '#/definitions/id'},
'graph': {'$ref': '#/definitions/graph'},
'dim': {'type': 'integer'},
'reshape_spec': {'$ref': '#/definitions/reshape'},
'architecture': {'$ref': '#/definitions/architecture'},
},
#'required': ['_class', '_id'],
'required': ['_class'],
'allOf': [
{
'if': {'properties': {'_class': {'enum': ['Sequential', 'Group']}}},
'then': {'required': ['blocks']},
'else': {'not': {'required': ['blocks']}},
},
{
'if': {'properties': {'_class': {'const': 'Group'}}},
'then': {'required': ['graph', 'input', 'output']},
'else': {'not': {'required': ['graph', 'input', 'output']}},
},
{
'if': {'properties': {'_class': {'const': 'Module'}}},
'then': {'required': ['_path']},
'else': {'not': {'required': ['_path', '_ext_vars', 'architecture']}},
},
{
'if': {'properties': {'_class': {'const': 'Sequential'}}},
'else': {'properties': {'blocks': {'items': {'required': ['_id']}}}}, # not working!
},
{
'if': {'properties': {'_class': {'const': 'Concatenate'}}},
'then': {'required': ['dim']},
},
{
'if': {'properties': {'_class': {'const': 'Reshape'}}},
'then': {'required': ['reshape_spec']},
'else': {'not': {'required': ['reshape_spec']}},
},
],
}
blocks_type = {
'type': 'array',
'minItems': 1,
'items': {'$ref': '#/definitions/block'},
}
inputs_outputs_type = {
'type': 'array',
'minItems': 1,
'items': {
'type': 'object',
'properties': {
'_id': {'$ref': '#/definitions/id'},
'_description': {'$ref': '#/definitions/description'},
'_shape': {'$ref': '#/definitions/dims'},
},
'required': ['_id', '_shape'],
'additionalProperties': False,
},
}
architecture_type = {
'type': 'object',
'properties': {
'_id': {'$ref': '#/definitions/id'},
'_description': {'$ref': '#/definitions/description'},
'blocks': {'$ref': '#/definitions/blocks'},
'graph': {'$ref': '#/definitions/graph'},
'inputs': {'$ref': '#/definitions/inputs_outputs'},
'outputs': {'$ref': '#/definitions/inputs_outputs'},
},
'required': ['_id', 'blocks', 'graph', 'inputs', 'outputs'],
'additionalProperties': False,
}
definitions = {
'id': id_type,
'description': description_type,
'dims': dims_type,
'dims_in': dims_in_type,
'shape': shape_type,
'graph': graph_type,
'block': block_type,
'blocks': blocks_type,
'path': path_type,
'inputs_outputs': inputs_outputs_type,
'architecture': architecture_type,
}
definitions.update(reshape_definitions)
narchi_schema = {
'$schema': 'http://json-schema.org/draft-07/schema#',
'$id': 'https://schema.omnius.com/json/narchi/1.0/schema.json',
'title': 'Neural Network Module Architecture Schema',
'$ref': '#/definitions/architecture',
'definitions': definitions,
}
block_definitions = deepcopy(definitions)
block_definitions['id'] = {'type': 'string', 'pattern': '^'+propagated_id_pattern+'$'}
block_definitions['graph']['items']['pattern'] = '^'+propagated_id_pattern+'( +-> +'+propagated_id_pattern+')+$'
block_schema = {
'$ref': '#/definitions/block',
'definitions': block_definitions,
}
propagated_definitions = deepcopy(block_definitions)
propagated_definitions['block']['required'] += ['_shape']
propagated_definitions['dims']['items']['oneOf'][1]['pattern'] = '^'+variable_pattern+'$'
propagated_definitions['architecture']['properties']['_shape'] = {'$ref': '#/definitions/shape'}
propagated_schema = deepcopy(narchi_schema)
propagated_schema['definitions'] = propagated_definitions
propagated_schema['title'] = 'Neural Network Module Propagated Architecture Schema'
del propagated_schema['$id']
mappings_schema = {
'type': 'object',
'minProperties': 1,
'patternProperties': {
'^'+id_pattern+'$': {
'type': 'object',
'properties': {
'class': {
'type': 'string',
'pattern': '^[A-Za-z_][0-9A-Za-z_.]*$'
},
'kwargs': {
'type': 'object',
'patternProperties': {
'^'+id_pattern+'$': {
'type': 'string',
'oneOf': [
{'pattern': '^'+id_pattern+'$'},
{'pattern': '^shape:in:(|-)[0-9]+$'},
{'pattern': '^const:str:[^:]+$'},
{'pattern': '^const:int:[0-9]+$'},
{'pattern': '^const:bool:(True|False)$'},
],
},
'^:skip:$': {
'type': 'string',
'pattern': '^'+id_pattern+'$',
},
},
'additionalProperties': False,
},
'required': ['class'],
'additionalProperties': False,
},
},
},
'additionalProperties': False,
}
block_validator = jsonvalidator(block_schema)
narchi_validator = jsonvalidator(narchi_schema)
propagated_validator = jsonvalidator(propagated_schema)
reshape_validator = jsonvalidator(reshape_schema)
mappings_validator = jsonvalidator(mappings_schema)
[docs]class SchemasEnum(enum.Enum):
"""Enum of the schemas defined in narchi."""
narchi = narchi_schema
"""Main schema which defines the general format for architecture files."""
propagated = propagated_schema
"""Schema for architectures in which the dimensions have been propagated."""
reshape = reshape_schema
"""Schema that defines the format to specify reshaping of tensors."""
block = block_schema
"""Schema for a single architecture block."""
mappings = mappings_schema
"""Schema for mappings between architectures and block implementations."""
schemas = {k: v.value for k, v in SchemasEnum.__dict__.items() if not k.startswith('_')}
schemas[None] = narchi_schema
[docs]def schema_as_str(schema: str = None):
"""Formats a schema as a pretty printed json string.
Args:
schema: The schema name to return among {'narchi', 'propagated', 'reshape', 'block', 'mappings'}.
Returns:
str: Pretty printed schema.
"""
return json.dumps(schemas[schema], indent=2, ensure_ascii=False)