sagemaker_schema_inference_from_dataframe_and_artifact_export.py

1import pandas as pd
2from sagemaker_schema_inference_artifacts import SchemaGenerator
3
4# Sample data
5data = {
6    "feature1": [1, 2, 3],
7    "feature2": ["A", "B", "C"],
8    "target": [0, 1, 0]
9}
10df = pd.DataFrame(data)
11
12# Initialize the SchemaGenerator
13generator = SchemaGenerator()
14
15# Generate the schema from the DataFrame
16schema = generator.generate_schema(df)
17
18# Save the schema to a local file
19output_path = "schema.json"
20generator.save_schema(schema, output_path)
21
22print(f"Schema successfully generated and saved to {output_path}")