<jupyter_start><jupyter_text>Useful functions (needs refactoring)<jupyter_code>import os, sys
root_dir, _ = os.path.split(os.getcwd())
script_dir = os.path.join(root_dir, 'scripts')
sys.path.append(script_dir)
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import tensorflow as tf
import numpy as np
import pprint
from hparams import hparams<jupyter_output><empty_output><jupyter_text>Dataset loading<jupyter_code>sound_feature_description = {
    "wav": tf.io.FixedLenFeature([], tf.string)
}

def _parse_sound_function(example_proto):
  x = tf.io.parse_single_example(example_proto, sound_feature_description)
  x['wav'] = tf.io.parse_tensor(x['wav'], out_type=hparams['ftype']) 
  return x

long_sound_feature_description = {
  "wav": tf.io.FixedLenFeature([], tf.string),
  "path": tf.io.FixedLenFeature([], tf.string),
  "number_of_slices": tf.io.FixedLenFeature([], tf.string)
}

def _parse_long_sound_function(example_proto):
  x = tf.io.parse_single_example(example_proto, long_sound_feature_description)
  x['wav'] = tf.io.parse_tensor(x['wav'], out_type=hparams['ftype'])
  x['path'] = tf.io.parse_tensor(x['path'], out_type=tf.string)
  x['number_of_slices'] = tf.io.parse_tensor(x['number_of_slices'], out_type=tf.int32)  
  return x
def load_single_file_tfrecords(record_file):
  raw_sound_dataset = tf.data.TFRecordDataset(record_file)
  parsed_sound_dataset = raw_sound_dataset.map(_parse_sound_function)
  return parsed_sound_dataset

def load_long_audio_tfrecords(record_file):
  raw_sound_dataset = tf.data.TFRecordDataset(record_file)
  parsed_sound_dataset = raw_sound_dataset.map(_parse_long_sound_function)
  return parsed_sound_dataset

def load_training_files_tfrecords(record_pattern):
  record_files = tf.data.TFRecordDataset.list_files(
    file_pattern=record_pattern)
  raw_sound_dataset = record_files.interleave(
    tf.data.TFRecordDataset,
    cycle_length=1,
    num_parallel_calls=tf.data.experimental.AUTOTUNE)
  parsed_sound_dataset = raw_sound_dataset.map(
    _parse_sound_function,
    num_parallel_calls=tf.data.experimental.AUTOTUNE)
  
  training_dataset = parsed_sound_dataset.shuffle(
    buffer_size=hparams['buffer_size']).batch(
    hparams['train_batch_size'],
    drop_remainder=True).prefetch(
    buffer_size=tf.data.experimental.AUTOTUNE)
  
  return training_dataset<jupyter_output><empty_output><jupyter_text>Optimizer compatibility with tf.float16 (Not working yet)<jupyter_code>def get_optimizer(hparams):
  """
  Return optimizer instance based on hparams
  
  Wrap the optimizer to avoid underflow if ftype=tf.float16
  """
  if hparams['optimizer'] == "Adam":
    if hparams['learning_rate_decay']:
      lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
        hparams['learning_rate'],
        decay_steps=200000,
        decay_rate=0.99,
        staircase=True)
      optimizer = tf.keras.optimizers.Adam(
      learning_rate=lr_schedule)
    else:
      optimizer = tf.keras.optimizers.Adam(
        learning_rate=hparams["learning_rate"])

  elif hparams['optimizer'] == "Adadelta":
    assert(hparams["learning_rate"] == 1.0), "Set learning_rate to 1.0"
    optimizer = tf.keras.optimizers.Adadelta(
      learning_rate=hparams['learning_rate'])
  else:
    raise ValueError("Supported Optimizer is either Adam or Adadelta")
    
  if hparams["mixed_precision"]:
    return tf.train.experimental.enable_mixed_precision_graph_rewrite(
      optimizer, "dynamic")
  else:
    return optimizer<jupyter_output><empty_output><jupyter_text>Discrete Logistic Mixture Helpers<jupyter_code>def log_sum_exp(x):
  """ numerically stable log_sum_exp implementation that prevents overflow """
  axis = len(x.shape)-1
  m = tf.reduce_max(x, axis)
  m2 = tf.reduce_max(x, axis, keepdims=True)
  return m + tf.math.log(tf.reduce_sum(tf.math.exp(x-m2), axis))

def log_prob_from_logits(x):
  """ numerically stable log_softmax implementation that prevents overflow """
  axis = len(x.shape)-1
  m = tf.reduce_max(x, axis, keepdims=True)
  return x - m - tf.math.log(tf.reduce_sum(tf.math.exp(x-m), axis, keepdims=True))<jupyter_output><empty_output><jupyter_text>Straight-through estimator<jupyter_code>@tf.custom_gradient
def round_ste(x):
  def grad(dy):
    return dy
  return tf.round(x), grad<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11764705882352941, 'fraction_numerical': 0.005835667600373483, 'mean_word_length': 5.6953125, 'pattern_counts': {'":': 6, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '444214', 'license': 'MIT', 'n_tokens_mistral': 1535, 'n_tokens_neox': 1519, 'n_words': 230, 'path': 'notebooks/training_utils.ipynb', 'repo_name': 'vatj/waveglow-craystack', 'size': '4284'}