Artifacts

def save_artifact(
    data,
    name: str,
    version: str = None,
    artifact_type = None,
    tags: list = None,
    extract_metadata: bool = True,
    include_visualizations: bool = True,
    user_metadata: dict = None,
    materializer: type = None,
    uri: str = None
):
    """
    Save an artifact to the artifact store.

    Parameters:
    - data: Data object to save
    - name: Artifact name
    - version: Artifact version (auto-generated if None)
    - artifact_type: Type of artifact (e.g., ArtifactType.MODEL, ArtifactType.DATA)
    - tags: List of tag names to attach
    - extract_metadata: Extract metadata automatically (default: True)
    - include_visualizations: Generate visualizations (default: True)
    - user_metadata: Custom metadata dict
    - materializer: Custom materializer class for serialization
    - uri: Optional URI to use for the artifact (advanced usage)

    Returns:
    ArtifactVersionResponse: Created artifact version object

    Example:
    ```python
    from zenml import save_artifact
    from zenml.enums import ArtifactType

    # Save a simple artifact
    artifact = save_artifact(
        data={"accuracy": 0.95, "loss": 0.05},
        name="model_metrics",
        tags=["evaluation", "production"]
    )
    print(f"Saved artifact: {artifact.id}")

    # Save with custom metadata and artifact type
    artifact = save_artifact(
        data=[1, 2, 3, 4, 5],
        name="training_data",
        version="v1.0",
        artifact_type=ArtifactType.DATA,
        user_metadata={"source": "database", "rows": 5}
    )
    ```
    """

from zenml import save_artifact

def load_artifact(
    name_or_id: str,
    version: str = None
):
    """
    Load an artifact from the artifact store.

    Parameters:
    - name_or_id: Artifact name or UUID
    - version: Artifact version (loads latest if None)

    Returns:
    Data object loaded from artifact store

    Example:
    ```python
    from zenml import load_artifact

    # Load latest version by name
    data = load_artifact("training_data")

    # Load specific version
    data = load_artifact("training_data", version="v1.0")

    # Load by UUID
    data = load_artifact("12345678-1234-1234-1234-123456789012")
    ```
    """

from zenml import load_artifact

def register_artifact(
    folder_or_file_uri: str,
    name: str,
    version: str = None,
    artifact_type=None,
    tags: list = None,
    has_custom_name: bool = True,
    artifact_metadata: dict = {}
):
    """
    Register existing data stored in the artifact store as a ZenML Artifact.

    The URI must point to a location within the active artifact store.

    Parameters:
    - folder_or_file_uri: Full URI within the artifact store to folder or file
    - name: Artifact name in ZenML
    - version: Artifact version (auto-incremented if None)
    - artifact_type: Type of artifact (defaults to 'data' if not given)
    - tags: List of tag names
    - has_custom_name: If artifact name is custom and should be listed in dashboard
    - artifact_metadata: Metadata dictionary to attach to the artifact version

    Returns:
    ArtifactVersionResponse: Registered artifact version object

    Raises:
    FileNotFoundError: If the folder URI is outside the artifact store bounds

    Example:
    ```python
    from zenml import register_artifact
    from zenml.enums import ArtifactType

    # Register artifact from artifact store
    artifact = register_artifact(
        folder_or_file_uri="s3://artifact-store/models/model.pkl",
        name="pretrained_model",
        version="v1.0",
        tags=["pretrained", "imported"],
        artifact_metadata={"source": "external"},
        artifact_type=ArtifactType.MODEL
    )

    # Register file in artifact store
    artifact = register_artifact(
        folder_or_file_uri="gs://artifact-store/data/data.csv",
        name="external_dataset"
    )
    ```
    """

from zenml import register_artifact

def log_artifact_metadata(
    metadata: dict,
    artifact_name: str = None,
    artifact_version: str = None
):
    """
    Log metadata for an artifact.

    Can be called within a step to attach metadata to output artifacts,
    or called outside step execution to attach metadata to any artifact.

    Parameters:
    - metadata: Metadata dict to log (keys must be strings)
    - artifact_name: Artifact name (uses current step output if None)
    - artifact_version: Artifact version

    Example:
    ```python
    from zenml import step, log_artifact_metadata

    @step
    def train_model(data: list) -> dict:
        model = {"weights": [0.1, 0.2], "accuracy": 0.95}

        # Log metadata for the output artifact
        log_artifact_metadata(
            metadata={
                "training_time": "120s",
                "epochs": 10,
                "optimizer": "adam"
            }
        )

        return model

    # Log metadata for existing artifact
    from zenml import log_artifact_metadata

    log_artifact_metadata(
        metadata={"reviewed": True, "reviewer": "data-team"},
        artifact_name="training_data",
        artifact_version="v1.0"
    )
    ```
    """

from zenml import log_artifact_metadata

class ArtifactType(str, Enum):
    """
    Types of artifacts.

    Values:
    - DATA: General data artifacts
    - MODEL: Model artifacts
    - DATA_ANALYSIS: Data analysis results
    - SERVICE: Service artifacts
    - STATISTICS: Statistical data
    - SCHEMA: Schema definitions
    - BASE: Base artifact type
    """
    DATA = "data"
    MODEL = "model"
    DATA_ANALYSIS = "data_analysis"
    SERVICE = "service"
    STATISTICS = "statistics"
    SCHEMA = "schema"
    BASE = "base"

from zenml.enums import ArtifactType

from zenml import save_artifact, load_artifact
from zenml.enums import ArtifactType

# Train a model
model_data = {
    "weights": [0.1, 0.2, 0.3],
    "bias": 0.5,
    "accuracy": 0.95
}

# Save model as artifact
saved = save_artifact(
    data=model_data,
    name="my_model",
    version="v1.0",
    artifact_type=ArtifactType.MODEL,
    tags=["production", "trained"],
    user_metadata={
        "framework": "custom",
        "training_time": "300s"
    }
)

print(f"Saved model artifact: {saved.id}")

# Later, load the model
loaded_model = load_artifact("my_model", version="v1.0")
print(f"Model accuracy: {loaded_model['accuracy']}")

from zenml import register_artifact
from zenml.enums import ArtifactType

# Register a pre-trained model in artifact store
model_artifact = register_artifact(
    folder_or_file_uri="s3://artifact-store/models/bert-base.bin",
    name="bert_base_pretrained",
    version="huggingface-v1",
    tags=["pretrained", "transformer", "nlp"],
    artifact_type=ArtifactType.MODEL,
    artifact_metadata={
        "source": "huggingface",
        "model_id": "bert-base-uncased",
        "parameters": "110M"
    }
)

# Register dataset in artifact store
data_artifact = register_artifact(
    folder_or_file_uri="gs://artifact-store/data/train.parquet",
    name="processed_training_data",
    tags=["training", "processed"],
    artifact_type=ArtifactType.DATA
)

from zenml import step, log_artifact_metadata
from typing import Tuple

@step
def preprocess_and_train(raw_data: list) -> Tuple[dict, dict]:
    """Preprocess data and train model."""

    # Preprocessing
    processed_data = [x * 2 for x in raw_data]

    # Log metadata for first output (processed_data)
    log_artifact_metadata(
        metadata={
            "preprocessing": "scaling",
            "scale_factor": 2,
            "rows": len(processed_data)
        }
    )

    # Training
    model = {"weights": [0.1, 0.2], "accuracy": 0.95}

    # Log metadata for second output (model)
    log_artifact_metadata(
        metadata={
            "training_epochs": 10,
            "optimizer": "adam",
            "learning_rate": 0.001
        }
    )

    return processed_data, model

from zenml import save_artifact, load_artifact, log_artifact_metadata

# Save multiple versions
for i in range(3):
    model = {"version": i, "accuracy": 0.9 + i * 0.01}
    save_artifact(
        data=model,
        name="evolving_model",
        version=f"v{i}.0",
        tags=["training"]
    )

# Load latest version
latest = load_artifact("evolving_model")
print(f"Latest version: {latest}")

# Load specific version
v1 = load_artifact("evolving_model", version="v1.0")
print(f"V1.0: {v1}")

# Add metadata to specific version
log_artifact_metadata(
    metadata={"status": "production", "deployed": True},
    artifact_name="evolving_model",
    artifact_version="v2.0"
)

from zenml import save_artifact, load_artifact
from zenml.materializers import CloudpickleMaterializer

# Custom object
class CustomModel:
    def __init__(self, weights):
        self.weights = weights

    def predict(self, x):
        return sum(w * x for w, x in zip(self.weights, x))

model = CustomModel([0.1, 0.2, 0.3])

# Save with custom materializer
save_artifact(
    data=model,
    name="custom_model",
    materializer=CloudpickleMaterializer
)

# Load
loaded = load_artifact("custom_model")
prediction = loaded.predict([1, 2, 3])
print(f"Prediction: {prediction}")

from zenml.client import Client

client = Client()

# Get artifact by name
artifact = client.get_artifact("training_data")
print(f"Artifact: {artifact.name}")

# List all versions
versions = client.list_artifact_versions(name="training_data")
for v in versions:
    print(f"Version {v.version}: {v.uri}")

# Get specific version
version = client.get_artifact_version("training_data", version="v1.0")
print(f"URI: {version.uri}")

# Update artifact metadata
client.update_artifact_version(
    name_or_id=version.id,
    add_tags=["validated"]
)

# Delete artifact version
client.delete_artifact_version(version.id)