# restart your notebook if prompted on Colab
try:
    import verta
except ImportError:
    !pip install verta

import os

# Ensure credentials are set up, if not, use below
# os.environ['VERTA_EMAIL'] = 
# os.environ['VERTA_DEV_KEY'] = 
# os.environ['VERTA_HOST'] = 

from verta import Client

client = Client(os.environ['VERTA_HOST'])

"""
Explanation: Creating a Standard Model on Verta
Within Verta, a "Model" can be any arbitrary function: a traditional ML model (e.g., sklearn, PyTorch, TF, etc); a function (e.g., squaring a number, making a DB function etc.); or a mixture of the above (e.g., pre-processing code, a DB call, and then a model application.) See more here.
This notebook provides an example of how to define a a Verta Standard Model by extending VertaModelBase. Verta also offers convenience functions for a number of common libraries including scikit-learn, PyTorch, Tensorflow as documented here.
0. Imports
End of explanation
"""


registered_model = client.get_or_create_registered_model(
    name="census", labels=["research-purpose", "team-a"])

from verta.registry import VertaModelBase

class MyModel(VertaModelBase):
    def __init__(self, artifacts):
        self.weights = json.load(open(artifacts["weights"]))
    
    def predict(self, input):
        res = []
        for row in input:
            res.append(row[0] * self.weights[0] + row[1] * self.weights[1])
        return res

from verta.environment import Python

model_version = registered_model.create_standard_model(
    model_cls=MyModel,
    artifacts = {"weights" : [5, 6]},
    environment=Python(requirements=["json"]),
    name="v0",
    labels=["prototype"],
)

"""
Explanation: 1. Register a model
End of explanation
"""


