Model

The Model class manages trained machine learning models in your apps.

Model Class

Model

Model(model: Any, label: str, metadata: Optional[dict] = None, name: Optional[str] = None, description: Optional[str] = None, version: Optional[int] = None, **kwargs: Any)

The model asset is a wrapper for machine learning models. Accessing attributes and member functions of the model asset also passes through access to the underlying machine learning model.

Parameters:

• model (Any) –

  The machine learning model that this asset keeps track of.

• label (str) –

  Label of :class:~virtualitics_sdk.assets.asset.Asset, see its documentation for more details.

• metadata (Optional[dict], default: None ) –

  Metadata of :class:~virtualitics_sdk.assets.asset.Asset, see its documentation for more details. Defaults to None.

• name (Optional[str], default: None ) –

  Name of :class:~virtualitics_sdk.assets.asset.Asset, see its documentation for more details. Defaults to None.

• description (Optional[str], default: None ) –

  Description of :class:~virtualitics_sdk.assets.asset.Asset, see its documentation for more details. Defaults to None.

• version (Optional[int], default: None ) –

  Version of :class:~virtualitics_sdk.assets.asset.Asset, see its documentation for more details. Defaults to 0. EXAMPLE: .. code-block:: python # Imports from virtualitics_sdk import Model from sklearn.linear_model import LogisticRegression . . . # Example usage X = np.random.rand(10).reshape(-1, 1) Y = np.random.randint(0, 2, size=(10, 1)) model = Model(LogisticRegression(), label="test", name="model") model.fit(X, Y) Additionally, for specific packages the model asset stores hyperparameter information and the time it took to run certain model functions. Currently supported packages are xgboost and sklearn.

update_model_notes

update_model_notes() -> None

This function updates the asset's stored metadata about the model.

get_all_metainfo

get_all_metainfo() -> Dict

Returns the dictionary containing the model metadata.

Returns:

• Dict –

  Model metadata, including hyperparameters and times taken for last function calls.

get_time

get_time(attr: str) -> Optional[str]

Returns the last time taken to run the function named attr. The function needs to have been called on the model asset rather than the underlying ML model in order for the time to be recorded. Returns None if no time was found.

Parameters:

• attr (str) –

  The name of the function to find the time taken.

Returns:

• Optional[str] –

  The time taken in seconds as a float. If no time is found, returns None instead.

get_recent_time

get_recent_time()

Returns the time taken for the most recent function which had its time recorded.

Returns:

• –

  The time taken in seconds as a float. If no time is found, returns None instead.

Basic Usage

from virtualitics_sdk import Model
from sklearn.ensemble import RandomForestClassifier

# Train a model
clf = RandomForestClassifier()
clf.fit(X_train, y_train)

# Wrap in Model asset
model = Model(
    name="Sales Predictor",
    model=clf
)

# Store for next step
self._outLink.trained_model = model

Using Stored Models

def run(self, flow_metadata):
    # Retrieve model
    model = self._inLink.trained_model

    # Access the underlying model object
    clf = model.model

    # Make predictions
    predictions = clf.predict(X_test)

    # Get model performance
    accuracy = clf.score(X_test, y_test)

Supported Model Types

The Model class works with various ML frameworks:

Scikit-learn

from sklearn.ensemble import RandomForestRegressor

model_obj = RandomForestRegressor()
model_obj.fit(X, y)

model = Model(name="RF Regressor", model=model_obj)

XGBoost

import xgboost as xgb

model_obj = xgb.XGBClassifier()
model_obj.fit(X, y)

model = Model(name="XGB Classifier", model=model_obj)

Custom Models

class CustomModel:
    def __init__(self):
        self.weights = None

    def fit(self, X, y):
        # Training logic
        pass

    def predict(self, X):
        # Prediction logic
        pass

custom_model = CustomModel()
custom_model.fit(X, y)

model = Model(name="Custom Model", model=custom_model)

Model Metadata

Store model information:

model = Model(
    name="Sales Predictor",
    model=clf,
    metadata={
        "algorithm": "Random Forest",
        "features": list(X.columns),
        "train_accuracy": train_score,
        "test_accuracy": test_score,
        "trained_at": datetime.now().isoformat(),
        "hyperparameters": {
            "n_estimators": 100,
            "max_depth": 10
        }
    }
)

# Access metadata
print(f"Model trained at: {model.metadata['trained_at']}")
print(f"Test accuracy: {model.metadata['test_accuracy']:.2%}")

Complete ML Workflow

class TrainModelStep(Step):
    def run(self, flow_metadata):
        # Get training data
        dataset = self._inLink.training_data
        df = dataset.data

        # Prepare features
        X = df.drop('target', axis=1)
        y = df['target']

        # Train model
        clf = RandomForestClassifier(n_estimators=100)
        clf.fit(X, y)

        # Calculate metrics
        train_score = clf.score(X, y)

        # Store model
        model = Model(
            name="Classifier",
            model=clf,
            metadata={
                "train_score": train_score,
                "features": list(X.columns)
            }
        )

        self._outLink.model = model

        return Page(...)


class PredictStep(Step):
    def run(self, flow_metadata):
        # Get model and new data
        model = self._inLink.model
        new_data = self._inLink.new_data

        # Make predictions
        predictions = model.model.predict(new_data.data)

        # Store predictions
        results = new_data.data.copy()
        results['prediction'] = predictions

        self._outLink.predictions = Dataset(
            name="Predictions",
            data=results
        )

        return Page(...)

Serialization

Models are automatically serialized using pickle/dill for storage:

# Model is automatically saved when stored in _outLink
self._outLink.model = model

# And automatically loaded when retrieved from _inLink
model = self._inLink.model

Best Practices

• Model Naming: Use descriptive names indicating model type and purpose
• Metadata: Store training metrics, feature names, and hyperparameters
• Validation: Validate model performance before storing
• Versioning: Include version info in metadata for model tracking
• Size: Very large models (>1GB) may have performance implications

See Also

• Assets Overview
• Dataset - Training data
• Schema - Data validation