wehub-resource-sync
130 lines
3.4 KiB
Python
130 lines
3.4 KiB
Python
"""Ondisk metadata of GraphBolt."""
|
|
|
|
from enum import Enum
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import pydantic
|
|
|
|
from ..internal_utils import version
|
|
|
|
|
|
__all__ = [
|
|
"OnDiskFeatureDataFormat",
|
|
"OnDiskTVTSetData",
|
|
"OnDiskTVTSet",
|
|
"OnDiskFeatureDataDomain",
|
|
"OnDiskFeatureData",
|
|
"OnDiskMetaData",
|
|
"OnDiskGraphTopologyType",
|
|
"OnDiskGraphTopology",
|
|
"OnDiskTaskData",
|
|
]
|
|
|
|
|
|
class ExtraMetaData(pydantic.BaseModel, extra="allow"):
|
|
"""Group extra fields into metadata. Internal use only."""
|
|
|
|
extra_fields: Optional[Dict[str, Any]] = {}
|
|
|
|
# As pydantic 2.0 has changed the API of validators, we need to use
|
|
# different validators for different versions to be compatible with
|
|
# previous versions.
|
|
if version.parse(pydantic.__version__) >= version.parse("2.0"):
|
|
|
|
@pydantic.model_validator(mode="before")
|
|
@classmethod
|
|
def build_extra(cls, values: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Build extra fields."""
|
|
for key in list(values.keys()):
|
|
if key not in cls.model_fields:
|
|
values["extra_fields"] = values.get("extra_fields", {})
|
|
values["extra_fields"][key] = values.pop(key)
|
|
return values
|
|
|
|
else:
|
|
|
|
@pydantic.root_validator(pre=True)
|
|
@classmethod
|
|
def build_extra(cls, values: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Build extra fields."""
|
|
for key in list(values.keys()):
|
|
if key not in cls.__fields__:
|
|
values["extra_fields"] = values.get("extra_fields", {})
|
|
values["extra_fields"][key] = values.pop(key)
|
|
return values
|
|
|
|
|
|
class OnDiskFeatureDataFormat(str, Enum):
|
|
"""Enum of data format."""
|
|
|
|
TORCH = "torch"
|
|
NUMPY = "numpy"
|
|
|
|
|
|
class OnDiskTVTSetData(pydantic.BaseModel):
|
|
"""Train-Validation-Test set data."""
|
|
|
|
name: Optional[str] = None
|
|
format: OnDiskFeatureDataFormat
|
|
in_memory: Optional[bool] = True
|
|
path: str
|
|
|
|
|
|
class OnDiskTVTSet(pydantic.BaseModel):
|
|
"""Train-Validation-Test set."""
|
|
|
|
type: Optional[str] = None
|
|
data: List[OnDiskTVTSetData]
|
|
|
|
|
|
class OnDiskFeatureDataDomain(str, Enum):
|
|
"""Enum of feature data domain."""
|
|
|
|
NODE = "node"
|
|
EDGE = "edge"
|
|
GRAPH = "graph"
|
|
|
|
|
|
class OnDiskFeatureData(ExtraMetaData):
|
|
r"""The description of an on-disk feature."""
|
|
domain: OnDiskFeatureDataDomain
|
|
type: Optional[str] = None
|
|
name: str
|
|
format: OnDiskFeatureDataFormat
|
|
path: str
|
|
in_memory: Optional[bool] = True
|
|
|
|
|
|
class OnDiskGraphTopologyType(str, Enum):
|
|
"""Enum of graph topology type."""
|
|
|
|
FUSED_CSC_SAMPLING = "FusedCSCSamplingGraph"
|
|
|
|
|
|
class OnDiskGraphTopology(pydantic.BaseModel):
|
|
"""The description of an on-disk graph topology."""
|
|
|
|
type: OnDiskGraphTopologyType
|
|
path: str
|
|
|
|
|
|
class OnDiskTaskData(ExtraMetaData):
|
|
"""Task specification in YAML."""
|
|
|
|
train_set: Optional[List[OnDiskTVTSet]] = []
|
|
validation_set: Optional[List[OnDiskTVTSet]] = []
|
|
test_set: Optional[List[OnDiskTVTSet]] = []
|
|
|
|
|
|
class OnDiskMetaData(pydantic.BaseModel):
|
|
"""Metadata specification in YAML.
|
|
|
|
As multiple node/edge types and multiple splits are supported, each TVT set
|
|
is a list of list of ``OnDiskTVTSet``.
|
|
"""
|
|
|
|
dataset_name: Optional[str] = None
|
|
graph_topology: Optional[OnDiskGraphTopology] = None
|
|
feature_data: Optional[List[OnDiskFeatureData]] = []
|
|
tasks: Optional[List[OnDiskTaskData]] = []
|