text: # coding=utf-8
  # Copyright 2021 The Google Research Authors.
  #
  # 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
  #
  #     http://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.
  
  # Lint as: python3
  r"""Construct a beam pipeline to map from audio to embeddings.
  
  This file has two modes:
  1) Map from tf.Examples of audio to tf.Examples of embeddings.
  2) Map from TFDS dataseet to tf.Examples of embeddings.
  
  It supports using a tf.hub module OR a TFLite model file to generate embeddings.
  TFLite file should have the `.tflite` extension.
  """
  
  import copy
  import numbers
  import os
  import random
  import typing
  from absl import logging
  import apache_beam as beam
  import librosa
  import numpy as np
  import tensorflow as tf
  import tensorflow_datasets as tfds
  import tensorflow_hub as hub
  from non_semantic_speech_benchmark import file_utils
  
  
  def _tfexample_audio_to_npfloat32(ex, audio_key):
    """Extract audio from tf.Example and convert it to np.float32."""
    audio_feats = ex.features.feature[audio_key]
    if audio_feats.int64_list.value:
      audio = np.array(audio_feats.int64_list.value)
      # Even though the data is in an int64 container, the data is actually int16.
      iinfo = np.iinfo(np.int16)
      assert np.logical_and(audio >= iinfo.min, audio <= iinfo.max).all(),\
          (np.min(audio), np.max(audio), iinfo.min, iinfo.max)
      audio = audio.astype(np.float32) / iinfo.max
    else:
      assert audio_feats.float_list.value
      audio = np.array(audio_feats.float_list.value, dtype=np.float32)
    return audio
  
  
  def _samples_to_embedding_tfhub(model_input, sample_rate, mod, output_key):
    """Run inference to map audio samples to an embedding."""
    tf_out = mod(
        tf.constant(model_input, tf.float32), tf.constant(sample_rate, tf.int32))
    return np.array(tf_out[output_key])
  
  
  def _build_tflite_interpreter(tflite_model_path):
    model_content = None
    with tf.io.gfile.GFile(tflite_model_path, 'rb') as model_file:
      model_content = model_file.read()
    interpreter = tf.lite.Interpreter(model_content=model_content)
    interpreter.allocate_tensors()
    return interpreter
  
  
  def _samples_to_embedding_tflite(model_input, sample_rate, interpreter,
                                   output_key):
    """Run TFLite inference to map audio samples to an embedding."""
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    # Resize TFLite input size based on length of sample.
    # Ideally, we should explore if we can use fixed-size input here, and
    # tile the sample to meet TFLite input size.
    interpreter.resize_tensor_input(input_details[0]['index'], model_input.shape)
    interpreter.allocate_tensors()
    interpreter.set_tensor(input_details[0]['index'], model_input)
    interpreter.set_tensor(input_details[1]['index'],
                           np.array(sample_rate).astype(np.int32))
  
    interpreter.invoke()
    embedding_2d = interpreter.get_tensor(
        output_details[int(output_key)]['index'])
    return np.array(embedding_2d, dtype=np.float32)
  
  
  @beam.typehints.with_input_types(typing.Tuple[str, typing.Any])
  @beam.typehints.with_output_types(typing.Tuple[str, typing.Any])
  class ComputeEmbeddingMapFn(beam.DoFn):
    """Computes an embedding (key, tf.Example) from audio (key, tf.Example)."""
  
    def __init__(self,
                 name,
                 module,
                 output_key,
                 audio_key,
                 sample_rate_key,
                 sample_rate,
                 average_over_time,
                 feature_fn=None):
      self._name = name
      # If TFLite should be used, `module` should point to a flatbuffer model.
      self._module = module
      self._use_tflite = self._module.endswith('.tflite')
      # For TFLite, `output_key` is the index of the embedding output from TFLite
      # model (Usually 0).
      self._output_key = output_key
      self._audio_key = audio_key
      self._sample_rate_key = sample_rate_key
      self._sample_rate = sample_rate
      self._average_over_time = average_over_time
      self._feature_fn = feature_fn
  
      # Only one of `sample_rate_key` and `sample_rate` should be not None.
      assert (self._sample_rate_key is None) ^ (self._sample_rate is None),\
          (self._sample_rate_key, self._sample_rate)
  
    def setup(self):
      if self._use_tflite:
        self.interpreter = _build_tflite_interpreter(self._module)
      else:
        self.module = hub.load(self._module)
  
    def process(self, k_v):
      k, ex = k_v
  
      # Read the input example audio and assert input format sanity.
      if self._audio_key not in ex.features.feature:
        raise ValueError(f'Audio key `{self._audio_key}` not found: '
                         f'{list(ex.features.feature.keys())}')
      audio = _tfexample_audio_to_npfloat32(ex, self._audio_key)
      if audio.size == 0:
        raise ValueError(f'No audio found: {self._audio_key}, {audio.size} {k}')
      beam.metrics.Metrics.distribution(
          'computed-embedding-audio', 'length').update(audio.size)
  
      # Read the sample rate, if a key to do so has been provided.
      if self._sample_rate_key:
        if self._sample_rate_key not in ex.features.feature:
          raise ValueError(f'Sample rate key not found: {self._sample_rate_key}')
        sample_rate = ex.features.feature[
            self._sample_rate_key].int64_list.value[0]
      else:
        if not self._sample_rate:
          raise ValueError('If `sample_rate_key` not provided, must provide '
                           '`sample_rate`.')
        sample_rate = self._sample_rate
      logging.info(
          'len(audio): %s / %s / %s', len(audio), sample_rate, self._name)
  
      # Resample, if necessary.
      if sample_rate != 16000:
        audio = librosa.core.resample(
            audio, orig_sr=sample_rate, target_sr=16000, res_type='kaiser_best')
        sample_rate = 16000
  
      # Convert audio to features, if required.
      if self._feature_fn:
        model_input = self._feature_fn(audio, sample_rate)
        if not isinstance(model_input, np.ndarray):
          raise ValueError(f'Expected ndarray, got {type(model_input)}')
        if model_input.dtype != np.float32:
          raise ValueError(f'Should be float32, was: {model_input.dtype}')
        logging.info('Audio shape is now: %s', model_input.shape)
      else:
        model_input = audio
  
      # Calculate the 2D embedding.
      if self._use_tflite:
        embedding_2d = _samples_to_embedding_tflite(
            audio, sample_rate, self.interpreter, self._output_key)
      else:
        embedding_2d = _samples_to_embedding_tfhub(model_input, sample_rate,
                                                   self.module, self._output_key)
      assert isinstance(embedding_2d, np.ndarray)
      assert embedding_2d.ndim == 2
      assert embedding_2d.dtype == np.float32
      beam.metrics.Metrics.counter('computed-embedding', self._name).inc()
      beam.metrics.Metrics.distribution(f'computed-embedding-{self._name}',
                                        'length').update(embedding_2d.shape[0])
  
      # Average over time, if required.
      if self._average_over_time:
        embedding = np.mean(embedding_2d, axis=0, keepdims=True)
      else:
        embedding = embedding_2d
  
      yield (k, embedding)
  
  
  def _add_embedding_to_tfexample(ex, embedding, name):
    """Add a 2D embedding to a tf.train.Example."""
    # Store the embedding 2D shape and store the 1D embedding. The original
    # embedding can be recovered with `emb.reshape(feature['shape'])`.
    f = ex.features.feature[f'{name}/shape']
    f.int64_list.value.extend(embedding.shape)
    f = ex.features.feature[name]
    f.float_list.value.extend(embedding.reshape([-1]))
    return ex
  
  
  def _add_embedding_column_map_fn(k_v, original_example_key,
                                   delete_audio_from_output, audio_key,
                                   label_key, speaker_id_key):
    """Combine a dictionary of named embeddings with a tf.train.Example."""
    k, v_dict = k_v
  
    assert original_example_key in v_dict, (original_example_key, v_dict.keys())
    ex_l = v_dict[original_example_key]
    assert len(ex_l) == 1, (len(ex_l), k_v[0], ex_l)
    ex = copy.deepcopy(ex_l[0])  # Beam does not allow modifying the input.
    assert isinstance(ex, tf.train.Example), type(ex)
  
    for name, embedding_l in v_dict.items():
      if name == original_example_key:
        continue
      assert len(embedding_l) == 1, embedding_l
      embedding = embedding_l[0]
      assert isinstance(embedding, np.ndarray)
      assert embedding.ndim == 2, embedding.ndim
  
      # Store the embedding 2D shape and store the 1D embedding. The original
      # embedding can be recovered with `emb.reshape(feature['shape'])`.
      ex = _add_embedding_to_tfexample(ex, embedding, f'embedding/{name}')
  
    if delete_audio_from_output:
      ex.features.feature.pop(audio_key, None)
  
    # Assert that the label is present. If it's a integer, convert it to bytes.
    assert label_key in ex.features.feature
    lbl_feat = ex.features.feature[label_key]
    if lbl_feat.int64_list.value:
      lbl_val_as_bytes = str(lbl_feat.int64_list.value[0]).encode('utf-8')
      ex.features.feature.pop(label_key, None)
      ex.features.feature[label_key].bytes_list.value.append(lbl_val_as_bytes)
  
    # If provided, assert that the speaker_id field is present, and of type
    # `bytes`.
    if speaker_id_key:
      feats = ex.features.feature
      assert speaker_id_key in feats, (speaker_id_key, feats.keys())
      assert feats[speaker_id_key].bytes_list.value, feats[speaker_id_key]
  
    return k, ex
  
  
  def _tfds_filenames(dataset_name, split_name, data_dir=None):
    """Returns filenames for a TFDS dataset."""
    data_dir = tfds.builder(dataset_name, data_dir=data_dir).data_dir
    return [os.path.join(data_dir, x) for x in
            tfds.builder(dataset_name).info.splits[split_name].filenames]
  
  
  def _tfds_sample_rate(dataset_name, data_dir=None):
    return tfds.builder(dataset_name, data_dir=data_dir).info.features[
        'audio'].sample_rate
  
  
  def read_input_glob_and_sample_rate_from_flags(
      input_glob_flag, sample_rate_flag, tfds_dataset_flag, output_filename_flag,
      tfds_data_dir_flag):
    """Read flags for input data and sample rate.
  
    Args:
      input_glob_flag: String flag. The input file glob.
      sample_rate_flag: String flag. The sample rate.
      tfds_dataset_flag: String flag. The TFDS dataset.
      output_filename_flag: String flag. The output filename.
      tfds_data_dir_flag: String flag. Optional location of local TFDS data.
  
    Returns:
      (input_filenames, output_filenames, sample_rate)
      `input_filenames` is a list of list of filenames. `output_filenames` is a
      list of the same length.
    """
    if input_glob_flag:
      assert file_utils.Glob(input_glob_flag), input_glob_flag
      assert not tfds_data_dir_flag
      input_filenames = [file_utils.Glob(input_glob_flag)]
      output_filenames = [output_filename_flag]
      sample_rate = sample_rate_flag
    else:
      assert tfds_dataset_flag
      dataset_name = tfds_dataset_flag
      # Download dataset, if necessary.
      tfds.load(dataset_name, data_dir=tfds_data_dir_flag)
      sample_rate = _tfds_sample_rate(dataset_name, tfds_data_dir_flag)
      assert sample_rate, sample_rate
  
      input_filenames = []
      output_filenames = []
      for split_name in ('train', 'validation', 'test'):
        input_filenames.append(
            _tfds_filenames(dataset_name, split_name, tfds_data_dir_flag))
        output_filenames.append(output_filename_flag + f'.{split_name}')
  
      logging.info('TFDS input filenames: %s', input_filenames)
      logging.info('sample rate: %s', sample_rate)
  
    if sample_rate:
      assert isinstance(sample_rate, numbers.Number)
  
    for filename_list in input_filenames:
      for filename in filename_list:
        assert tf.io.gfile.exists(filename), filename
    assert len(input_filenames) == len(output_filenames)
  
    return input_filenames, output_filenames, sample_rate
  
  
  def validate_inputs(
      input_filenames_list, output_filenames, embedding_modules, embedding_names,
      module_output_keys):
    """Validate inputs and input flags."""
    for filename_list in input_filenames_list:
      for filename in filename_list:
        assert tf.io.gfile.exists(filename), filename
    assert len(input_filenames_list) == len(output_filenames)
  
    # Make sure output files don't already exist.
    for output_filename in output_filenames:
      assert not file_utils.Glob(f'{output_filename}*'), output_filename
  
    # Lengths of flag lists must be the same.
    assert len(embedding_names) == len(embedding_modules),\
           (embedding_names, embedding_modules)
    assert len(embedding_modules) == len(module_output_keys),\
           (len(embedding_modules), len(module_output_keys))
    # Shortnames must be unique.
    assert len(set(embedding_names)) == len(embedding_names), embedding_names
  
    # Create output directory if it doesn't already exist.
    for output_filename in output_filenames:
      output_dir = output_filename.rsplit('/', 1)[0]
      file_utils.MaybeMakeDirs(output_dir)
  
  
  def _read_from_tfrecord(root, input_filenames, suffix):
    """Reads from a Python list of TFRecord files."""
    assert isinstance(input_filenames, list), input_filenames
    return (root
            | f'MakeFilenames{suffix}' >> beam.Create(input_filenames)
            | f'ReadTFRecords{suffix}' >> beam.io.tfrecordio.ReadAllFromTFRecord(
                coder=beam.coders.ProtoCoder(tf.train.Example))
            | f'AddKeys{suffix}' >> beam.Map(
                lambda x: (str(random.getrandbits(128)), x)))
  
  
  
  
  def _write_to_tfrecord(combined_tbl, output_filename, suffix):
    _ = (combined_tbl
         | f'RemoveKey{suffix}' >> beam.Map(lambda k_v: k_v[1])
         | f'Write{suffix}' >> beam.io.WriteToTFRecord(
             output_filename, coder=beam.coders.ProtoCoder(tf.train.Example)))
  
  
  
  
  # Possible input formats. If you want to read from a different input format,
  # add your read function here. Function should take (root, input_filenames) and
  # map to input_examples.
  reader_functions = {
      'tfrecord': _read_from_tfrecord
  }
  
  
  # Write output to disk.
  writer_functions = {
      'tfrecord': _write_to_tfrecord,
  }
  
  
  def make_beam_pipeline(
      root, input_filenames, sample_rate, debug, embedding_names,
      embedding_modules, module_output_keys, audio_key, sample_rate_key,
      label_key, speaker_id_key, average_over_time, delete_audio_from_output,
      output_filename, input_format='tfrecord', output_format='tfrecord',
      suffix='Main'):
    """Construct beam pipeline for mapping from audio to embeddings.
  
    Args:
      root: The beam root node.
      input_filenames: Python list. List of input files.
      sample_rate: Python int, or `None`. The sample rate for all embeddings,
        or `None` if this is a TFDS dataset, or if each example has its own sample
        rate.
      debug: Python bool. Whether to operate in debug mode.
      embedding_names: Python list of embeddings.
      embedding_modules: Python list of TF-Hub modules.
      module_output_keys: Python list of strings, names of output modules.
      audio_key: Python string, the key of the audio.
      sample_rate_key: Python string or `None`, the key for.
      label_key: Python string. Field for label.
      speaker_id_key: Python string or `None`. Key for speaker ID, or `None`.
      average_over_time: Python bool. If `True`, average over the time axis.
      delete_audio_from_output: Python bool. Whether to remove audio fromm
        outputs.
      output_filename: Python string. Output filename.
      input_format: Python string. Must correspond to a function in
        `reader_functions`.
      output_format: Python string. Must correspond to a function
        `writer_functions`.
      suffix: Python string. Suffix to stage names to make them unique.
    """
    tf_examples_key_ = 'tf_examples'
    assert tf_examples_key_ not in embedding_names
    s = suffix  # for code brevity.
  
    # Read from input.
    input_examples = reader_functions[input_format](root, input_filenames, s)
  
    # In debug mode, take one input example.
    if debug:
      input_examples = (
          input_examples
          | f'TakeOne{s}' >> beam.transforms.combiners.Sample.FixedSizeGlobally(1)
          # Sampling generates lists, so flatten back into one collection.
          | f'DebugFlatten{s}' >> beam.FlatMap(lambda x: x))
  
    # Compute all the embeddings simultaneously.
    embedding_tables = {}
    for name, mod, out_key in zip(
        embedding_names, embedding_modules, module_output_keys):
      logging.info('Adding signal: %s %s, %s', name, mod, out_key)
      tbl = input_examples | f'ComputeEmbedding-{name}-{s}' >> beam.ParDo(
          ComputeEmbeddingMapFn(
              name=name,
              module=mod,
              output_key=out_key,
              audio_key=audio_key,
              sample_rate_key=sample_rate_key,
              sample_rate=sample_rate,
              average_over_time=average_over_time))
      embedding_tables[name] = tbl
    assert tf_examples_key_ not in embedding_tables
    embedding_tables[tf_examples_key_] = input_examples
    logging.info('embedding_tables: %s', embedding_tables)
  
    # Combine embeddings and tf.train.Example, using the common key.
    combined_tbl = (
        embedding_tables
        | f'CombineEmbeddingTables-{s}' >> beam.CoGroupByKey()
        | f'AddEmbeddings-{s}' >> beam.Map(
            _add_embedding_column_map_fn,
            original_example_key=tf_examples_key_,
            delete_audio_from_output=delete_audio_from_output,
            audio_key=audio_key,
            label_key=label_key,
            speaker_id_key=speaker_id_key))
  
    output_filename = f'{output_filename}@*'
    logging.info('Writing to %s', output_filename)
    writer_functions[output_format](combined_tbl, output_filename, s)
  