title: Saving a 'fine-tuned' bert model
body: <p>I am trying to save a fine tuned bert model. I have ran the code correctly - it works fine, and in the ipython console I am able to call getPrediction and have it result the result. </p>
  
  <p>I have my weight files saved (highest being model.ckpt-333.data-00000-of-00001</p>
  
  <p>I have no idea how I would go about saving the model to be reuseable. </p>
  
  <p>I am using bert-tensorflow. </p>
  
  <pre><code>import json
  
  import pandas as pd
  import tensorflow as tf
  import tensorflow_hub as hub
  from datetime import datetime
  
  
  from sklearn.model_selection import train_test_split
  import os
  
  print("tensorflow version : ", tf.__version__)
  print("tensorflow_hub version : ", hub.__version__)
  
  
  #Importing BERT modules
  import bert
  from bert import run_classifier
  from bert import optimization
  from bert import tokenization
  
  #set output directory of the model
  OUTPUT_DIR = 'model'
  
  #@markdown Whether or not to clear/delete the directory and create a new one
  DO_DELETE = False #@param {type:"boolean"}
  
  if DO_DELETE:
    try:
      tf.gfile.DeleteRecursively(OUTPUT_DIR)
    except:
      pass
  
  tf.io.gfile.makedirs(OUTPUT_DIR)
  print('***** Model output directory: {} *****'.format(OUTPUT_DIR))
  
  
  ### Load the data
  data = pd.read_csv("data/bbc-text.csv")
  
  data.columns = ['category', 'text']
  print('*****Data Loaded: {} *****'.format(data.head()))
  
  #check to see if any null values are present.
  print('*****Empty Data: {} *****'.format(data[data.isnull().any(axis=1)]))
  
  #encode category variable into numeric
  data.category = pd.Categorical(data.category)
  data['code'] = data.category.cat.codes
  
  from sklearn.model_selection import train_test_split
  
  train, test = train_test_split(data, test_size=0.2, random_state=200)
  
  ## 2 -- Data Visualisation
  
  print(data.code.unique())
  
  import matplotlib.pyplot as plt
  
  train['code'].value_counts().plot(kind = 'bar')
  DATA_COLUMN = 'text'
  LABEL_COLUMN = 'code'
  label_list = [0, 1, 2, 3, 4]
  plt.show()
  
  ## 2 -- Data Preprocessing
  
  train_InputExamples = train.apply(lambda x: bert.run_classifier.InputExample(guid=None,
                                                                     text_a = x[DATA_COLUMN],
                                                                     text_b = None,
                                                                     label = x[LABEL_COLUMN]), axis = 1)
  
  test_InputExamples = test.apply(lambda x: bert.run_classifier.InputExample(guid=None,
                                                                     text_a = x[DATA_COLUMN],
                                                                     text_b = None,
                                                                     label = x[LABEL_COLUMN]), axis = 1)
  
  # This is a path to an uncased (all lowercase) version of BERT
  BERT_MODEL_HUB = "https://tfhub.dev/google/bert_uncased_L-12_H-768_A-12/1"
  
  def create_tokenizer_from_hub_module():
    """Get the vocab file and casing info from the Hub module."""
    with tf.Graph().as_default():
      bert_module = hub.Module(BERT_MODEL_HUB)
      tokenization_info = bert_module(signature="tokenization_info", as_dict=True)
      with tf.compat.v1.Session() as sess:
        vocab_file, do_lower_case = sess.run([tokenization_info["vocab_file"],
                                              tokenization_info["do_lower_case"]])
  
    return bert.tokenization.FullTokenizer(
        vocab_file=vocab_file, do_lower_case=do_lower_case)
  
  tokenizer = create_tokenizer_from_hub_module()
  
  # We'll set sequences to be at most 128 tokens long.
  MAX_SEQ_LENGTH = 128
  
  # Convert our train and validation features to InputFeatures that BERT understands.
  train_features = bert.run_classifier.convert_examples_to_features(train_InputExamples, label_list, MAX_SEQ_LENGTH, tokenizer)
  
  test_features = bert.run_classifier.convert_examples_to_features(test_InputExamples, label_list, MAX_SEQ_LENGTH, tokenizer)
  
  
  #Example on first observation in the training set
  print("Example of train[0] as a training set")
  print("Sentence : ", train_InputExamples.iloc[0].text_a)
  print("-"*30)
  print("Tokens : ", tokenizer.tokenize(train_InputExamples.iloc[0].text_a))
  print("-"*30)
  print("Input IDs : ", train_features[0].input_ids)
  print("-"*30)
  print("Input Masks : ", train_features[0].input_mask)
  print("-"*30)
  print("Segment IDs : ", train_features[0].segment_ids)
  
  
  ## 3. Creating a Multiclass Classifier
  def create_model(is_predicting, input_ids, input_mask, segment_ids, labels,
                   num_labels):
  
    bert_module = hub.Module(
        BERT_MODEL_HUB,
        trainable=True)
    bert_inputs = dict(
        input_ids=input_ids,
        input_mask=input_mask,
        segment_ids=segment_ids)
    bert_outputs = bert_module(
        inputs=bert_inputs,
        signature="tokens",
        as_dict=True)
  
    # Use "pooled_output" for classification tasks on an entire sentence.
    # Use "sequence_outputs" for token-level output.
    output_layer = bert_outputs["pooled_output"]
  
    hidden_size = output_layer.shape[-1].value
  
    # Create our own layer to tune for politeness data.
    output_weights = tf.compat.v1.get_variable(
        "output_weights", [num_labels, hidden_size],
        initializer=tf.truncated_normal_initializer(stddev=0.02))
  
    output_bias = tf.compat.v1.get_variable(
        "output_bias", [num_labels], initializer=tf.zeros_initializer())
  
    with tf.compat.v1.variable_scope("loss"):
  
      # Dropout helps prevent overfitting
      output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)
  
      logits = tf.matmul(output_layer, output_weights, transpose_b=True)
      logits = tf.nn.bias_add(logits, output_bias)
      log_probs = tf.nn.log_softmax(logits, axis=-1)
  
      # Convert labels into one-hot encoding
      one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)
  
      predicted_labels = tf.squeeze(tf.argmax(log_probs, axis=-1, output_type=tf.int32))
      # If we're predicting, we want predicted labels and the probabiltiies.
      if is_predicting:
        return (predicted_labels, log_probs)
  
      # If we're train/eval, compute loss between predicted and actual label
      per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)
      loss = tf.reduce_mean(per_example_loss)
      return (loss, predicted_labels, log_probs)
  
  #A function that adapts our model to work for training, evaluation, and prediction.
  
  # model_fn_builder actually creates our model function
  # using the passed parameters for num_labels, learning_rate, etc.
  def model_fn_builder(num_labels, learning_rate, num_train_steps,
                       num_warmup_steps):
    """Returns `model_fn` closure for TPUEstimator."""
    def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
      """The `model_fn` for TPUEstimator."""
  
      input_ids = features["input_ids"]
      input_mask = features["input_mask"]
      segment_ids = features["segment_ids"]
      label_ids = features["label_ids"]
  
      is_predicting = (mode == tf.estimator.ModeKeys.PREDICT)
  
      # TRAIN and EVAL
      if not is_predicting:
  
        (loss, predicted_labels, log_probs) = create_model(
          is_predicting, input_ids, input_mask, segment_ids, label_ids, num_labels)
  
        train_op = bert.optimization.create_optimizer(
            loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu=False)
  
        # Calculate evaluation metrics.
        def metric_fn(label_ids, predicted_labels):
          accuracy = tf.compat.v1.metrics.accuracy(label_ids, predicted_labels)
          true_pos = tf.compat.v1.metrics.true_positives(
              label_ids,
              predicted_labels)
          true_neg = tf.compat.v1.metrics.true_negatives(
              label_ids,
              predicted_labels)
          false_pos = tf.compat.v1.metrics.false_positives(
              label_ids,
              predicted_labels)
          false_neg = tf.compat.v1.metrics.false_negatives(
              label_ids,
              predicted_labels)
  
          return {
              "eval_accuracy": accuracy,
              "true_positives": true_pos,
              "true_negatives": true_neg,
              "false_positives": false_pos,
              "false_negatives": false_neg
              }
  
        eval_metrics = metric_fn(label_ids, predicted_labels)
  
        if mode == tf.estimator.ModeKeys.TRAIN:
          return tf.estimator.EstimatorSpec(mode=mode,
            loss=loss,
            train_op=train_op)
        else:
            return tf.estimator.EstimatorSpec(mode=mode,
              loss=loss,
              eval_metric_ops=eval_metrics)
      else:
        (predicted_labels, log_probs) = create_model(
          is_predicting, input_ids, input_mask, segment_ids, label_ids, num_labels)
  
        predictions = {
            'probabilities': log_probs,
            'labels': predicted_labels
        }
        return tf.estimator.EstimatorSpec(mode, predictions=predictions)
  
    # Return the actual model function in the closure
    return model_fn
  
  # Compute train and warmup steps from batch size
  # These hyperparameters are copied from this colab notebook (https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/tools/colab/bert_finetuning_with_cloud_tpus.ipynb)
  BATCH_SIZE = 16
  LEARNING_RATE = 2e-5
  NUM_TRAIN_EPOCHS = 3.0
  # Warmup is a period of time where the learning rate is small and gradually increases--usually helps training.
  WARMUP_PROPORTION = 0.1
  # Model configs
  SAVE_CHECKPOINTS_STEPS = 300
  SAVE_SUMMARY_STEPS = 100
  
  # Compute train and warmup steps from batch size
  num_train_steps = int(len(train_features) / BATCH_SIZE * NUM_TRAIN_EPOCHS)
  num_warmup_steps = int(num_train_steps * WARMUP_PROPORTION)
  
  # Specify output directory and number of checkpoint steps to save
  run_config = tf.estimator.RunConfig(
      model_dir=OUTPUT_DIR,
      save_summary_steps=SAVE_SUMMARY_STEPS,
      save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS)
  
  # Specify output directory and number of checkpoint steps to save
  run_config = tf.estimator.RunConfig(
      model_dir=OUTPUT_DIR,
      save_summary_steps=SAVE_SUMMARY_STEPS,
      save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS)
  
  
  #Initializing the model and the estimator
  model_fn = model_fn_builder(
    num_labels=len(label_list),
    learning_rate=LEARNING_RATE,
    num_train_steps=num_train_steps,
    num_warmup_steps=num_warmup_steps)
  
  estimator = tf.estimator.Estimator(
    model_fn=model_fn,
    config=run_config,
    params={"batch_size": BATCH_SIZE})
  
  # Create an input function for training. drop_remainder = True for using TPUs.
  train_input_fn = bert.run_classifier.input_fn_builder(
      features=train_features,
      seq_length=MAX_SEQ_LENGTH,
      is_training=True,
      drop_remainder=False)
  
  # Create an input function for validating. drop_remainder = True for using TPUs.
  test_input_fn = run_classifier.input_fn_builder(
      features=test_features,
      seq_length=MAX_SEQ_LENGTH,
      is_training=False,
      drop_remainder=False)
  
  
  # #Training the model
  print(f'Beginning Training!')
  current_time = datetime.now()
  estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)
  print("Training took time ", datetime.now() - current_time)
  
  #Evaluating the model with Validation set
  accuracy = estimator.evaluate(input_fn=test_input_fn, steps=None)
  
  
  # A method to get predictions
  def getPrediction(in_sentences):
      # A list to map the actual labels to the predictions
      labels = ["business", "entertainment", "politics", "sports", "tech"]
  
  
      # Transforming the test data into BERT accepted form
      input_examples = [run_classifier.InputExample(guid="", text_a=x, text_b=None, label=0) for x in in_sentences]
  
      # Creating input features for Test data
      input_features = run_classifier.convert_examples_to_features(input_examples, label_list, MAX_SEQ_LENGTH, tokenizer)
  
      # Predicting the classes
      predict_input_fn = run_classifier.input_fn_builder(features=input_features, seq_length=MAX_SEQ_LENGTH,
                                                         is_training=False, drop_remainder=False)
      predictions = estimator.predict(predict_input_fn)
      return [(sentence, prediction['probabilities'], prediction['labels'], labels[prediction['labels']]) for
              sentence, prediction in zip(in_sentences, predictions)]
  pred_sentences = list(test['text'])
  
  predictions = getPrediction(pred_sentences)
  
  enc_labels = []
  act_labels = []
  for i in range(len(predictions)):
    enc_labels.append(predictions[i][2])
    act_labels.append(predictions[i][3])
  
  pd.DataFrame(enc_labels, columns = ['category']).to_excel('data/submission_bert.xlsx', index = False)
  
  ## Random tester
  #Classifying random sentences
  tests = getPrediction(['Mr.Modi is the Indian Prime Minister',
                         'Gaming machines are powered by efficient micro processores and GPUs',
                         'That HBO TV series is really good',
                         'A trillion dollar economy '
                         ])
  
  </code></pre>