!pip3 uninstall -y google-cloud-aiplatform
!pip3 install google-cloud-aiplatform
import IPython

app = IPython.Application.instance()
app.kernel.do_shutdown(True)

import sys

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

"""
Explanation: Feedback or issues?
For any feedback or questions, please open an issue.
Vertex SDK for Python: Custom Tabular Model Training Example
To use this Colaboratory notebook, you copy the notebook to your own Google Drive and open it with Colaboratory (or Colab). You can run each step, or cell, and see its results. To run a cell, use Shift+Enter. Colab automatically displays the return value of the last line in each cell. For more information about running notebooks in Colab, see the Colab welcome page.
This notebook demonstrate how to create a custom model based on a tabular dataset. It will require you provide a bucket where the dataset CSV will be stored.
Note: you may incur charges for training, prediction,  storage or usage of other GCP products in connection with testing this SDK.
Install Vertex SDK for Python, Authenticate, and upload of a Dataset to your GCS bucket
After the SDK installation the kernel will be automatically restarted. You may see this error message Your session crashed for an unknown reason which is normal.
End of explanation
"""


MY_PROJECT = "YOUR PROJECT"
MY_STAGING_BUCKET = "gs://YOUR BUCKET"  # bucket should be in same region as Vertex AI

"""
Explanation: Enter your project and GCS bucket
Enter your Project Id in the cell below. Then run the cell to make sure the Cloud SDK uses the right project for all the commands in this notebook.
End of explanation
"""


!wget https://storage.googleapis.com/download.tensorflow.org/data/abalone_train.csv
!gsutil cp abalone_train.csv {MY_STAGING_BUCKET}/data/

gcs_csv_path = f"{MY_STAGING_BUCKET}/data/abalone_train.csv"

"""
Explanation: The dataset we are using is the Abalone Dataset. For more information about this dataset please visit: https://archive.ics.uci.edu/ml/datasets/abalone
End of explanation
"""


from google.cloud import aiplatform

aiplatform.init(project=MY_PROJECT, staging_bucket=MY_STAGING_BUCKET)

"""
Explanation: Initialize Vertex SDK for Python
Initialize the client for Vertex AI
End of explanation
"""


ds = aiplatform.TabularDataset.create(display_name="abalone", gcs_source=[gcs_csv_path])

ds.resource_name

"""
Explanation: Create a Managed Tabular Dataset from CSV
A Managed dataset can be used to create an AutoML model or a custom model.
End of explanation
"""


%%writefile training_script.py

import pandas as pd
import os
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers


# uncomment and bump up replica_count for distributed training
# strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
# tf.distribute.experimental_set_strategy(strategy)

col_names = ["Length", "Diameter", "Height", "Whole weight", "Shucked weight", "Viscera weight", "Shell weight", "Age"]
target = "Age"

def aip_data_to_dataframe(wild_card_path):
    return pd.concat([pd.read_csv(fp.numpy().decode(), names=col_names)
                      for fp in tf.data.Dataset.list_files([wild_card_path])])

def get_features_and_labels(df):
    return df.drop(target, axis=1).values, df[target].values

def data_prep(wild_card_path):
    return get_features_and_labels(aip_data_to_dataframe(wild_card_path))


model = tf.keras.Sequential([layers.Dense(64), layers.Dense(1)])
model.compile(loss='mse', optimizer='adam')

model.fit(*data_prep(os.environ["AIP_TRAINING_DATA_URI"]),
          epochs=10 ,
          validation_data=data_prep(os.environ["AIP_VALIDATION_DATA_URI"]))
print(model.evaluate(*data_prep(os.environ["AIP_TEST_DATA_URI"])))

# save as Vertex AI Managed model
tf.saved_model.save(model, os.environ["AIP_MODEL_DIR"])

"""
Explanation: Write Training Script

Write this cell as a file which will be used for custom training.
End of explanation
"""


job = aiplatform.CustomTrainingJob(
    display_name="train-abalone-dist-1-replica",
    script_path="training_script.py",
    container_uri="gcr.io/cloud-aiplatform/training/tf-cpu.2-2:latest",
    requirements=["gcsfs==0.7.1"],
    model_serving_container_image_uri="gcr.io/cloud-aiplatform/prediction/tf2-cpu.2-2:latest",
)
model = job.run(ds, replica_count=1, model_display_name="abalone-model")

"""
Explanation: Launch a Training Job to Create a Model
Once we have defined your training script, we will create a model.
End of explanation
"""


endpoint = model.deploy(machine_type="n1-standard-4")

"""
Explanation: Deploy Model
End of explanation
"""


prediction = endpoint.predict(
    [
        [0.435, 0.335, 0.11, 0.33399999999999996, 0.1355, 0.0775, 0.0965],
        [0.585, 0.45, 0.125, 0.874, 0.3545, 0.2075, 0.225],
    ]
)
prediction

"""
Explanation: Predict on Endpoint
End of explanation
"""


