!git clone https://github.com/google-research/google-research.git

import sys
import os
import tarfile
import urllib
import zipfile
sys.path.append('./google-research')

"""
Explanation: Copyright 2019 Google LLC
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
End of explanation
"""


# TF streaming
from kws_streaming.models import models
from kws_streaming.models import utils
from kws_streaming.layers.modes import Modes

import tensorflow as tf
import numpy as np
import tensorflow.compat.v1 as tf1
import logging
from kws_streaming.models import model_flags
from kws_streaming.models import model_params
from kws_streaming.train import test
from kws_streaming.train import train
from kws_streaming import data
tf1.disable_eager_execution()

config = tf1.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf1.Session(config=config)

# general imports
import matplotlib.pyplot as plt
import os
import json
import numpy as np
import scipy as scipy
import scipy.io.wavfile as wav
import scipy.signal

tf.__version__

tf1.reset_default_graph()
sess = tf1.Session()
tf1.keras.backend.set_session(sess)
tf1.keras.backend.set_learning_phase(0)

"""
Explanation: Example of model training
Below steps are taken from model_train_eval - it has more tests in streaming, non streaming, quantized and non qunatized models with TF and TFLite.
Imports
End of explanation
"""


# set PATH to data sets (for example to speech commands V2):
# it can be downloaded from
# https://storage.googleapis.com/download.tensorflow.org/data/speech_commands_v0.02.tar.gz
# if you already run "00_check-data.ipynb" then folder "data2" should be located in the current dir
current_dir = os.getcwd()
DATA_PATH = os.path.join(current_dir, "data2/")

def waveread_as_pcm16(filename):
  """Read in audio data from a wav file.  Return d, sr."""
  samplerate, wave_data = wav.read(filename)
  # Read in wav file.
  return wave_data, samplerate

def wavread_as_float(filename, target_sample_rate=16000):
  """Read in audio data from a wav file.  Return d, sr."""
  wave_data, samplerate = waveread_as_pcm16(filename)
  desired_length = int(
          round(float(len(wave_data)) / samplerate * target_sample_rate))
  wave_data = scipy.signal.resample(wave_data, desired_length)

  # Normalize short ints to floats in range [-1..1).
  data = np.array(wave_data, np.float32) / 32768.0
  return data, target_sample_rate

# Set path to wav file to visualize it
wav_file = os.path.join(DATA_PATH, "left/012187a4_nohash_0.wav")

# read audio file
wav_data, samplerate = wavread_as_float(wav_file)

assert samplerate == 16000

plt.plot(wav_data)

"""
Explanation: Set path to data
End of explanation
"""


# select model name should be one of
model_params.HOTWORD_MODEL_PARAMS.keys()

# This notebook is configured to work with 'ds_tc_resnet' and 'svdf'.
MODEL_NAME = 'ds_tc_resnet'
# MODEL_NAME = 'svdf'
MODELS_PATH = os.path.join(current_dir, "models")
MODEL_PATH = os.path.join(MODELS_PATH, MODEL_NAME + "/")
MODEL_PATH

# delete previously trained model with its folder and create a new one:
os.makedirs(MODEL_PATH)

# get toy model settings
FLAGS = model_params.HOTWORD_MODEL_PARAMS[MODEL_NAME]

# set path to data and model (where model will be stored)
FLAGS.data_dir = DATA_PATH
FLAGS.train_dir = MODEL_PATH

# set speech feature extractor properties
FLAGS.mel_upper_edge_hertz = 7600
FLAGS.window_size_ms = 30.0
FLAGS.window_stride_ms = 10.0
FLAGS.mel_num_bins = 80
FLAGS.dct_num_features = 40
FLAGS.feature_type = 'mfcc_tf'
FLAGS.preprocess = 'raw'

# for numerical correctness of streaming and non streaming models set it to 1
# but for real use case streaming set it to 0
FLAGS.causal_data_frame_padding = 0

FLAGS.use_tf_fft = True
FLAGS.mel_non_zero_only = not FLAGS.use_tf_fft

# set training settings
FLAGS.train = 1
# reduced number of training steps for test only
# so model accuracy will be low,
# to improve accuracy set how_many_training_steps = '40000,40000,20000,20000'
FLAGS.how_many_training_steps = '400,400,400,400'
FLAGS.learning_rate = '0.001,0.0005,0.0001,0.00002'
FLAGS.lr_schedule = 'linear'
FLAGS.verbosity = logging.INFO

# data augmentation parameters
FLAGS.resample = 0.15
FLAGS.time_shift_ms = 100
FLAGS.use_spec_augment = 1
FLAGS.time_masks_number = 2
FLAGS.time_mask_max_size = 25
FLAGS.frequency_masks_number = 2
FLAGS.frequency_mask_max_size = 7
FLAGS.pick_deterministically = 1

FLAGS.model_name = MODEL_NAME

# model parameters are different for every model
if MODEL_NAME == 'svdf':
  FLAGS.model_name = MODEL_NAME
  FLAGS.svdf_memory_size = "4,10,10,10,10,10"
  FLAGS.svdf_units1 = "16,32,32,32,64,128"
  FLAGS.svdf_act = "'relu','relu','relu','relu','relu','relu'"
  FLAGS.svdf_units2 = "40,40,64,64,64,-1"
  FLAGS.svdf_dropout = "0.0,0.0,0.0,0.0,0.0,0.0"
  FLAGS.svdf_pad = 0
  FLAGS.dropout1 = 0.0
  FLAGS.units2 = ''
  FLAGS.act2 = ''
elif MODEL_NAME == 'ds_tc_resnet':
  # it is an example of model streaming with strided convolution, strided pooling and dilated convolution
  FLAGS.activation = 'relu'
  FLAGS.dropout = 0.0
  FLAGS.ds_filters = '128, 64, 64, 64, 128, 128'
  FLAGS.ds_filter_separable = '1, 1, 1, 1, 1, 1'
  FLAGS.ds_repeat = '1, 1, 1, 1, 1, 1'
  FLAGS.ds_residual = '0, 1, 1, 1, 0, 0' # residual can not be applied with stride
#   FLAGS.ds_kernel_size = '11, 5, 15, 7, 29, 1'
  FLAGS.ds_kernel_size = '11, 5, 15, 17, 15, 1'
  FLAGS.ds_dilation = '1, 1, 1, 1, 2, 1'
  FLAGS.ds_stride = '1, 1, 1, 1, 1, 1'
  FLAGS.ds_pool = '1, 2, 1, 1, 1, 1'
  # model should be causal, so that we can covert it to streaming mode
  # if model is non causal then all non causal components should use Delay layer
  FLAGS.ds_padding = "'causal', 'causal', 'causal', 'causal', 'causal', 'causal'"
else:
  raise ValueError("set parameters for other models")

FLAGS.clip_duration_ms = 1000  # standard audio file in this data set has 1 sec length
FLAGS.batch_size = 100

flags = model_flags.update_flags(FLAGS)

with open(os.path.join(flags.train_dir, 'flags.json'), 'wt') as f:
  json.dump(flags.__dict__, f)

# visualize a model
model_non_stream_batch = models.MODELS[flags.model_name](flags)
tf.keras.utils.plot_model(
    model_non_stream_batch,
    show_shapes=True,
    show_layer_names=True,
    expand_nested=True)

model_non_stream_batch.summary()

"""
Explanation: Set path to a model with config
End of explanation
"""


# Model training
train.train(flags)

"""
Explanation: Model training
End of explanation
"""


folder_name = 'tf'
test.tf_non_stream_model_accuracy(flags, folder_name)

"""
Explanation: Run model evaluation
TF Run non streaming inference
End of explanation
"""


