try:
 from google.colab import auth
 auth.authenticate_user()
except:
 pass

"""
Explanation: Serving Text Classification Using Keras on AI Platform
Overview
This notebook illustrates the new feature of serving custom model prediction code on AI Platform. It allows us to execute arbitrary python pre-processing code prior to invoking a model, as well as post-processing on the produced predictions.
This is all done server-side so that the client can pass data directly to AI Platform Serving in the unprocessed state.
We will take advantage of this for text classification because it involves pre-processing that is not easily accomplished using native TensorFlow. Instead we will execute the the non TensorFlow pre-processing via python code on the server side.
We implement our Text Classifier using Keras. Keras is a high-level API for building and training deep learning models. tf.keras is TensorFlow’s implementation of this API. To learn more about building machine learning models in Keras more generally, read TensorFlow's Keras tutorials.
Dataset
Hacker News is one of many public datasets available in BigQuery. This dataset includes titles of articles from several data sources. For the following tutorial, we extracted the titles that belong to either GitHub, The New York Times, or TechCrunch, and saved them as CSV files in a publicly shared Cloud Storage bucket at the following location: gs://cloud-training-demos/blogs/CMLE_custom_prediction
Objective
The goal of this tutorial is to:
1. Process the data for text classification.
2. Train a Keras Text Classifier (locally).
3. Deploy the Keras Text Classifier, along with the preprocessing artifacts, to AI Platform Serving, using the Custom Online Prediction code.
This tutorial focuses more on using this model with AI Platform Serving than on the design of the text classification model itself. For more details about text classification, please refer to Google developer's Guide to Text Classification. 
Costs
This tutorial uses billable components of Google Cloud Platform (GCP):
1. AI Platform Serving (Cloud Machine Learning Engine)
2. Cloud Storage
Learn about AI Platform Serving pricing and Cloud Storage pricing, and use the Pricing Calculator to generate a cost estimate based on your projected usage.
<a href="https://colab.research.google.com/github/GoogleCloudPlatform/cloudml-samples/blob/main/notebooks/keras/Text%20Classification%20Using%20Keras%20and%20CMLE.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
If you are using AI Platform Notebooks, your environment is already
authenticated. Skip this step.
End of explanation
"""


%load_ext autoreload
%autoreload 2

!pip install tensorflow==1.13.1

import tensorflow as tf
print(tf.__version__) 

import os

PROJECT='' # SET TO YOUR GCP PROJECT NAME
BUCKET='' # SET TO YOUR GCS BUCKET NAME
ROOT='keras_text_classification'
MODEL_DIR=os.path.join(ROOT,'models')
PACKAGES_DIR=os.path.join(ROOT,'packages')

!gcloud config set project {PROJECT}

# Delete any previous artefacts from Google Cloud Storage
!gsutil rm -r gs://{BUCKET}/{ROOT}

"""
Explanation: Setup
End of explanation
"""


# Download from Google Cloud Storage
%%bash
gsutil cp gs://cloud-training-demos/blogs/CMLE_custom_prediction/keras_text_pre_processing/train.tsv .
gsutil cp gs://cloud-training-demos/blogs/CMLE_custom_prediction/keras_text_pre_processing/eval.tsv .

!head eval.tsv

"""
Explanation: Download and Explore Data
End of explanation
"""


%%writefile preprocess.py

from tensorflow.python.keras.preprocessing import sequence
from tensorflow.keras.preprocessing import text

class TextPreprocessor(object):
  def __init__(self, vocab_size, max_sequence_length):
    self._vocab_size = vocab_size
    self._max_sequence_length = max_sequence_length
    self._tokenizer = None

  def fit(self, text_list):        
    # Create vocabulary from input corpus.
    tokenizer = text.Tokenizer(num_words=self._vocab_size)
    tokenizer.fit_on_texts(text_list)
    self._tokenizer = tokenizer

  def transform(self, text_list):        
    # Transform text to sequence of integers
    text_sequence = self._tokenizer.texts_to_sequences(text_list)

    # Fix sequence length to max value. Sequences shorter than the length are
    # padded in the beginning and sequences longer are truncated
    # at the beginning.
    padded_text_sequence = sequence.pad_sequences(
      text_sequence, maxlen=self._max_sequence_length)
    return padded_text_sequence

"""
Explanation: Preprocessing
Pre-processing class to be used in both training and serving
End of explanation
"""


from preprocess import TextPreprocessor

processor = TextPreprocessor(5, 5)
processor.fit(['hello machine learning'])
processor.transform(['hello machine learning'])

"""
Explanation: Test Prepocessing Locally
End of explanation
"""


CLASSES = {'github': 0, 'nytimes': 1, 'techcrunch': 2}  # label-to-int mapping
VOCAB_SIZE = 20000  # Limit on the number vocabulary size used for tokenization
MAX_SEQUENCE_LENGTH = 50  # Sentences will be truncated/padded to this length

"""
Explanation: Model Creation
Metadata
End of explanation
"""


import pandas as pd
import numpy as np
from preprocess import TextPreprocessor

def load_data(train_data_path, eval_data_path):
    # Parse CSV using pandas
    column_names = ('label', 'text')
    df_train = pd.read_csv(
      train_data_path, names=column_names, sep='\t').sample(frac=1)
    df_eval = pd.read_csv(
      eval_data_path, names=column_names, sep='\t')

    return ((list(df_train['text']), np.array(df_train['label'].map(CLASSES))),
            (list(df_eval['text']), np.array(df_eval['label'].map(CLASSES))))


((train_texts, train_labels), (eval_texts, eval_labels)) = load_data(
       'train.tsv', 'eval.tsv')

# Create vocabulary from training corpus.
processor = TextPreprocessor(VOCAB_SIZE, MAX_SEQUENCE_LENGTH)
processor.fit(train_texts)

# Preprocess the data
train_texts_vectorized = processor.transform(train_texts)
eval_texts_vectorized = processor.transform(eval_texts)

"""
Explanation: Prepare data for training and evaluation
End of explanation
"""


import pickle

with open('./processor_state.pkl', 'wb') as f:
  pickle.dump(processor, f)

"""
Explanation: Save pre-processing object
We need to save this so the same tokenizer used at training can be used to pre-process during serving
End of explanation
"""


import tensorflow as tf

from tensorflow.keras import models
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Dropout
from tensorflow.keras.layers import Embedding
from tensorflow.keras.layers import Conv1D
from tensorflow.keras.layers import MaxPooling1D
from tensorflow.keras.layers import GlobalAveragePooling1D


tf.logging.set_verbosity(tf.logging.INFO)

#keras model

def create_model(vocab_size, embedding_dim, 
  filters, kernel_size, dropout_rate, pool_size):
  
  model = models.Sequential()
  model.add(Embedding(input_dim=vocab_size,
                          output_dim=embedding_dim,
                          input_length=MAX_SEQUENCE_LENGTH))

  model.add(Dropout(rate=dropout_rate))
  model.add(Conv1D(filters=filters,
                            kernel_size=kernel_size,
                            activation='relu',
                            bias_initializer='random_uniform',
                            padding='same'))

  model.add(MaxPooling1D(pool_size=pool_size))
  model.add(Conv1D(filters=filters * 2,
                            kernel_size=kernel_size,
                            activation='relu',
                            bias_initializer='random_uniform',
                            padding='same'))
  model.add(GlobalAveragePooling1D())
  model.add(Dropout(rate=dropout_rate))
  model.add(Dense(len(CLASSES), activation='softmax'))

  return model


"""
Explanation: Build the model
End of explanation
"""


LEARNING_RATE=.001
EMBEDDING_DIM=200
FILTERS=64
KERNEL_SIZE=3
DROPOUT_RATE=0.2
POOL_SIZE=3

NUM_EPOCH=1
BATCH_SIZE=128

model = create_model(VOCAB_SIZE, EMBEDDING_DIM, 
  FILTERS, KERNEL_SIZE, DROPOUT_RATE,POOL_SIZE)

# Compile model with learning parameters.
optimizer = tf.keras.optimizers.Adam(lr=LEARNING_RATE)
model.compile(
  optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['acc'])

#keras train
model.fit(
  train_texts_vectorized, train_labels, epochs=NUM_EPOCH, batch_size=BATCH_SIZE)
print('Eval loss/accuracy:{}'.format(
  model.evaluate(eval_texts_vectorized, eval_labels, batch_size=BATCH_SIZE)))

#save model
model.save('keras_saved_model.h5')

"""
Explanation: Train and save the model
End of explanation
"""


!gsutil cp keras_saved_model.h5 gs://{BUCKET}/{MODEL_DIR}/
!gsutil cp processor_state.pkl gs://{BUCKET}/{MODEL_DIR}/

"""
Explanation: Custom Model Prediction Preparation
Copy model and pre-processing object to GCS
End of explanation
"""


%%writefile model_prediction.py

import os
import pickle
import numpy as np


class CustomModelPrediction(object):

  def __init__(self, model, processor):
    self._model = model
    self._processor = processor

  def _postprocess(self, predictions):
    labels = ['github', 'nytimes', 'techcrunch']
    label_indexes = [np.argmax(prediction) for prediction in predictions]
    return [labels[label_index] for label_index in label_indexes]


  def predict(self, instances, **kwargs):
    preprocessed_data = self._processor.transform(instances)
    predictions =  self._model.predict(preprocessed_data)
    labels = self._postprocess(predictions)
    return labels


  @classmethod
  def from_path(cls, model_dir):
    import tensorflow.keras as keras
    model = keras.models.load_model(
      os.path.join(model_dir,'keras_saved_model.h5'))
    with open(os.path.join(model_dir, 'processor_state.pkl'), 'rb') as f:
      processor = pickle.load(f)

    return cls(model, processor)
 

"""
Explanation: Define Model Class
End of explanation
"""


# Headlines for Predictions

techcrunch=[
  'Uber shuts down self-driving trucks unit',
  'Grover raises €37M Series A to offer latest tech products as a subscription',
  'Tech companies can now bid on the Pentagon’s $10B cloud contract'
]
nytimes=[
  '‘Lopping,’ ‘Tips’ and the ‘Z-List’: Bias Lawsuit Explores Harvard’s Admissions',
  'A $3B Plan to Turn Hoover Dam into a Giant Battery',
  'A MeToo Reckoning in China’s Workplace Amid Wave of Accusations'
]
github=[
  'Show HN: Moon – 3kb JavaScript UI compiler',
  'Show HN: Hello, a CLI tool for managing social media',
  'Firefox Nightly added support for time-travel debugging'
]
requests = (techcrunch+nytimes+github)

from model_prediction import CustomModelPrediction

classifier = CustomModelPrediction.from_path('.')
results = classifier.predict(requests)
results

"""
Explanation: Test Model Class Locally
End of explanation
"""


%%writefile setup.py

from setuptools import setup

setup(
  name="my_package",
  version="0.1",
  include_package_data=True,
  scripts=["preprocess.py", "model_prediction.py"]
)

!python setup.py sdist
!gsutil cp ./dist/my_package-0.1.tar.gz gs://{BUCKET}/{PACKAGES_DIR}/my_package-0.1.tar.gz

"""
Explanation: Package up files and copy to GCS
End of explanation
"""


MODEL_NAME='keras_text_classification'
VERSION_NAME='v201903'
RUNTIME_VERSION='1.13'
REGION='us-central1'

!gcloud ai-platform models create {MODEL_NAME} --regions {REGION}

!gcloud ai-platform versions delete {VERSION_NAME} --model {MODEL_NAME} --quiet # run if version already created

!gcloud beta ai-platform versions create {VERSION_NAME} --model {MODEL_NAME} \
--origin=gs://{BUCKET}/{MODEL_DIR}/ \
--python-version=3.5 \
--runtime-version={RUNTIME_VERSION} \
--framework='SCIKIT_LEARN' \
--package-uris=gs://{BUCKET}/{PACKAGES_DIR}/my_package-0.1.tar.gz \
--prediction-class=model_prediction.CustomModelPrediction

"""
Explanation: Model Deployment to CMLE
End of explanation
"""


from googleapiclient import discovery
from oauth2client.client import GoogleCredentials
import json

# JSON format the requests
request_data = {'instances': requests}

# Authenticate and call CMLE prediction API 
credentials = GoogleCredentials.get_application_default()
api = discovery.build(
  'ml', 'v1', credentials=credentials,
  discoveryServiceUrl='https://storage.googleapis.com/cloud-ml/discovery/ml_v1_discovery.json')

parent = 'projects/{}/models/{}/versions/{}'.format(PROJECT, MODEL_NAME, VERSION_NAME)
print("Model full name: {}".format(parent))
response = api.projects().predict(body=request_data, name=parent).execute()

print(response['predictions'])

"""
Explanation: Online Predictions from CMLE
End of explanation
"""


