sagemaker_schema_inference_artifacts_build_serialize_deserialize.py

1import json
2from sagemaker_schema_inference_artifacts.schema_builder import SchemaBuilder
3from sagemaker_schema_inference_artifacts.serializers import JSONSerializer, JSONDeserializer
4
5# 1. Define sample input and output data
6sample_input = {"features": [5.1, 3.5, 1.4, 0.2]}
7sample_output = {"prediction": 0}
8
9# 2. Initialize the SchemaBuilder with sample data
10# This automatically infers the schema from the provided samples
11schema_builder = SchemaBuilder(
12    sample_input=sample_input,
13    sample_output=sample_output
14)
15
16# 3. Generate the schema artifacts
17schema_artifacts = schema_builder.build()
18
19# 4. Serialize the schema to a JSON string for storage/transmission
20serialized_schema = schema_artifacts.serialize(serializer=JSONSerializer())
21print("Serialized Schema:")
22print(serialized_schema)
23
24# 5. Deserialize the schema back into a Python object
25deserialized_schema = schema_artifacts.deserialize(
26    serialized_schema, 
27    deserializer=JSONDeserializer()
28)
29
30# 6. Access the inferred input and output schemas
31print("\nInferred Input Schema:")
32print(json.dumps(deserialized_schema.input_schema, indent=2))
33
34print("\nInferred Output Schema:")
35print(json.dumps(deserialized_schema.output_schema, indent=2))