mirror of https://github.com/kubeflow/examples.git
288 lines
9.4 KiB
Python
288 lines
9.4 KiB
Python
"""Train the github-issue-summarization model
|
|
train.py trains the github-issue-summarization model.
|
|
|
|
It reads the input data from GCS in a zip file format.
|
|
--input_data_gcs_bucket and --input_data_gcs_path specify
|
|
the location of input data.
|
|
|
|
It write the model back to GCS.
|
|
--output_model_gcs_bucket and --output_model_gcs_path specify
|
|
the location of output.
|
|
|
|
It also has parameters which control the training like
|
|
--learning_rate and --sample_size
|
|
|
|
"""
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import re
|
|
import zipfile
|
|
|
|
from google.cloud import storage # pylint: disable=no-name-in-module
|
|
import dill as dpickle
|
|
import numpy as np
|
|
import pandas as pd
|
|
from keras import optimizers
|
|
from keras.layers import GRU, BatchNormalization, Dense, Embedding, Input
|
|
from keras.models import Model
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
from ktext.preprocess import processor
|
|
from seq2seq_utils import load_encoder_inputs, load_text_processor
|
|
|
|
GCS_REGEX = re.compile("gs://([^/]*)(/.*)?")
|
|
|
|
|
|
def split_gcs_uri(gcs_uri):
|
|
"""Split a GCS URI into bucket and path."""
|
|
m = GCS_REGEX.match(gcs_uri)
|
|
bucket = m.group(1)
|
|
path = ""
|
|
if m.group(2):
|
|
path = m.group(2).lstrip("/")
|
|
return bucket, path
|
|
|
|
|
|
def main(): # pylint: disable=too-many-statements
|
|
# Parsing flags.
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--sample_size", type=int, default=2000000)
|
|
parser.add_argument("--learning_rate", default="0.001")
|
|
|
|
parser.add_argument(
|
|
"--input_data",
|
|
type=str,
|
|
default="",
|
|
help="The input location. Can be a GCS or local file path.")
|
|
|
|
# TODO(jlewi): The following arguments are deprecated; just
|
|
# use input_data. We should remove them as soon as all call sites
|
|
# are updated.
|
|
parser.add_argument(
|
|
"--input_data_gcs_bucket", type=str, default="kubeflow-examples")
|
|
parser.add_argument(
|
|
"--input_data_gcs_path",
|
|
type=str,
|
|
default="github-issue-summarization-data/github-issues.zip")
|
|
|
|
parser.add_argument(
|
|
"--output_model",
|
|
type=str,
|
|
default="",
|
|
help="The output location for the model GCS or local file path.")
|
|
|
|
parser.add_argument("--output_model_gcs_bucket", type=str, default="")
|
|
parser.add_argument(
|
|
"--output_model_gcs_path",
|
|
type=str,
|
|
default="github-issue-summarization-data")
|
|
|
|
parser.add_argument(
|
|
"--output_body_preprocessor_dpkl",
|
|
type=str,
|
|
default="body_pp.dpkl")
|
|
parser.add_argument(
|
|
"--output_title_preprocessor_dpkl",
|
|
type=str,
|
|
default="title_pp.dpkl")
|
|
parser.add_argument(
|
|
"--output_train_title_vecs_npy", type=str, default="train_title_vecs.npy")
|
|
parser.add_argument(
|
|
"--output_train_body_vecs_npy", type=str, default="train_body_vecs.npy")
|
|
parser.add_argument(
|
|
"--output_model_h5", type=str, default="seq2seq_model_tutorial.h5")
|
|
|
|
args = parser.parse_args()
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format=('%(levelname)s|%(asctime)s'
|
|
'|%(pathname)s|%(lineno)d| %(message)s'),
|
|
datefmt='%Y-%m-%dT%H:%M:%S',
|
|
)
|
|
logging.getLogger().setLevel(logging.INFO)
|
|
logging.info(args)
|
|
|
|
learning_rate = float(args.learning_rate)
|
|
|
|
pd.set_option('display.max_colwidth', 500)
|
|
|
|
# For backwords compatibility
|
|
input_data_gcs_bucket = None
|
|
input_data_gcs_path = None
|
|
|
|
if not args.input_data:
|
|
# Since input_data isn't set fall back on old arguments.
|
|
input_data_gcs_bucket = args.input_data_gcs_bucket
|
|
input_data_gcs_path = args.input_data_gcs_path
|
|
else:
|
|
if args.input_data.startswith('gs://'):
|
|
input_data_gcs_bucket, input_data_gcs_path = split_gcs_uri(
|
|
args.input_data)
|
|
|
|
if input_data_gcs_bucket:
|
|
logging.info("Download bucket %s object %s.", input_data_gcs_bucket,
|
|
input_data_gcs_path)
|
|
bucket = storage.Bucket(storage.Client(), input_data_gcs_bucket)
|
|
args.input_data = 'github-issues.zip'
|
|
storage.Blob(input_data_gcs_path, bucket).download_to_filename(
|
|
args.input_data)
|
|
|
|
ext = os.path.splitext(args.input_data)[-1]
|
|
if ext.lower() == '.zip':
|
|
zip_ref = zipfile.ZipFile(args.input_data, 'r')
|
|
zip_ref.extractall('.')
|
|
zip_ref.close()
|
|
# TODO(jlewi): Hardcoding the file in the Archive to use is brittle.
|
|
# We should probably just require the input to be a CSV file.
|
|
csv_file = 'github_issues.csv'
|
|
else:
|
|
csv_file = args.input_data
|
|
|
|
# Read in data sample 2M rows (for speed of tutorial)
|
|
traindf, testdf = train_test_split(
|
|
pd.read_csv(csv_file).sample(n=args.sample_size), test_size=.10)
|
|
|
|
# Print stats about the shape of the data.
|
|
logging.info('Train: %d rows %d columns', traindf.shape[0], traindf.shape[1])
|
|
logging.info('Test: %d rows %d columns', testdf.shape[0], testdf.shape[1])
|
|
|
|
train_body_raw = traindf.body.tolist()
|
|
train_title_raw = traindf.issue_title.tolist()
|
|
|
|
# Clean, tokenize, and apply padding / truncating such that each document
|
|
# length = 70. Also, retain only the top 8,000 words in the vocabulary and set
|
|
# the remaining words to 1 which will become common index for rare words.
|
|
body_pp = processor(keep_n=8000, padding_maxlen=70)
|
|
train_body_vecs = body_pp.fit_transform(train_body_raw)
|
|
|
|
logging.info('Example original body: %s', train_body_raw[0])
|
|
logging.info('Example body after pre-processing: %s', train_body_vecs[0])
|
|
|
|
# Instantiate a text processor for the titles, with some different parameters.
|
|
title_pp = processor(
|
|
append_indicators=True, keep_n=4500, padding_maxlen=12, padding='post')
|
|
|
|
# process the title data
|
|
train_title_vecs = title_pp.fit_transform(train_title_raw)
|
|
|
|
logging.info('Example original title: %s', train_title_raw[0])
|
|
logging.info('Example title after pre-processing: %s', train_title_vecs[0])
|
|
|
|
# Save the preprocessor.
|
|
with open(args.output_body_preprocessor_dpkl, 'wb') as f:
|
|
dpickle.dump(body_pp, f)
|
|
|
|
with open(args.output_title_preprocessor_dpkl, 'wb') as f:
|
|
dpickle.dump(title_pp, f)
|
|
|
|
# Save the processed data.
|
|
np.save(args.output_train_title_vecs_npy, train_title_vecs)
|
|
np.save(args.output_train_body_vecs_npy, train_body_vecs)
|
|
|
|
_, doc_length = load_encoder_inputs(args.output_train_body_vecs_npy)
|
|
|
|
num_encoder_tokens, body_pp = load_text_processor(
|
|
args.output_body_preprocessor_dpkl)
|
|
num_decoder_tokens, title_pp = load_text_processor(
|
|
args.output_title_preprocessor_dpkl)
|
|
|
|
# Arbitrarly set latent dimension for embedding and hidden units
|
|
latent_dim = 300
|
|
|
|
###############
|
|
# Encoder Model.
|
|
###############
|
|
encoder_inputs = Input(shape=(doc_length,), name='Encoder-Input')
|
|
|
|
# Word embeding for encoder (ex: Issue Body)
|
|
x = Embedding(
|
|
num_encoder_tokens, latent_dim, name='Body-Word-Embedding',
|
|
mask_zero=False)(encoder_inputs)
|
|
x = BatchNormalization(name='Encoder-Batchnorm-1')(x)
|
|
|
|
# We do not need the `encoder_output` just the hidden state.
|
|
_, state_h = GRU(latent_dim, return_state=True, name='Encoder-Last-GRU')(x)
|
|
|
|
# Encapsulate the encoder as a separate entity so we can just
|
|
# encode without decoding if we want to.
|
|
encoder_model = Model(
|
|
inputs=encoder_inputs, outputs=state_h, name='Encoder-Model')
|
|
|
|
seq2seq_encoder_out = encoder_model(encoder_inputs)
|
|
|
|
################
|
|
# Decoder Model.
|
|
################
|
|
decoder_inputs = Input(
|
|
shape=(None,), name='Decoder-Input') # for teacher forcing
|
|
|
|
# Word Embedding For Decoder (ex: Issue Titles)
|
|
dec_emb = Embedding(
|
|
num_decoder_tokens,
|
|
latent_dim,
|
|
name='Decoder-Word-Embedding',
|
|
mask_zero=False)(decoder_inputs)
|
|
dec_bn = BatchNormalization(name='Decoder-Batchnorm-1')(dec_emb)
|
|
|
|
# Set up the decoder, using `decoder_state_input` as initial state.
|
|
decoder_gru = GRU(
|
|
latent_dim, return_state=True, return_sequences=True, name='Decoder-GRU')
|
|
decoder_gru_output, _ = decoder_gru(dec_bn, initial_state=seq2seq_encoder_out)
|
|
x = BatchNormalization(name='Decoder-Batchnorm-2')(decoder_gru_output)
|
|
|
|
# Dense layer for prediction
|
|
decoder_dense = Dense(
|
|
num_decoder_tokens, activation='softmax', name='Final-Output-Dense')
|
|
decoder_outputs = decoder_dense(x)
|
|
|
|
################
|
|
# Seq2Seq Model.
|
|
################
|
|
|
|
seq2seq_Model = Model([encoder_inputs, decoder_inputs], decoder_outputs)
|
|
|
|
seq2seq_Model.compile(
|
|
optimizer=optimizers.Nadam(lr=learning_rate),
|
|
loss='sparse_categorical_crossentropy')
|
|
|
|
seq2seq_Model.summary()
|
|
|
|
#############
|
|
# Save model.
|
|
#############
|
|
seq2seq_Model.save(args.output_model_h5)
|
|
|
|
######################
|
|
# Upload model to GCS.
|
|
######################
|
|
# For backwords compatibility
|
|
output_model_gcs_bucket = None
|
|
output_model_gcs_path = None
|
|
|
|
if not args.output_model:
|
|
# Since input_data isn't set fall back on old arguments.
|
|
output_model_gcs_bucket = args.output_model_gcs_bucket
|
|
output_model_gcs_path = args.output_model_gcs_path
|
|
else:
|
|
if args.output_model.startswith('gs://'):
|
|
output_model_gcs_bucket, output_model_gcs_path = split_gcs_uri(
|
|
args.output_model)
|
|
|
|
if output_model_gcs_bucket:
|
|
logging.info("Uploading model files to bucket %s path %s.",
|
|
output_model_gcs_bucket, output_model_gcs_path)
|
|
bucket = storage.Bucket(storage.Client(), output_model_gcs_bucket)
|
|
storage.Blob(
|
|
output_model_gcs_path + "/" + args.output_model_h5, bucket).upload_from_filename(
|
|
args.output_model_h5)
|
|
storage.Blob(output_model_gcs_path + "/" + args.output_body_preprocessor_dpkl,
|
|
bucket).upload_from_filename(args.output_body_preprocessor_dpkl)
|
|
storage.Blob(output_model_gcs_path + "/" + args.output_title_preprocessor_dpkl,
|
|
bucket).upload_from_filename(args.output_title_preprocessor_dpkl)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|