def test_compute_embedding_map_fn_tflite(self, average_over_time, sample_rate_key, sample_rate): # Establish required key names. audio_key = 'audio_key' # Construct the tf.train.Example test data. ex = tf.train.Example() ex.features.feature[audio_key].float_list.value.extend( np.zeros(2000, np.float32)) if sample_rate_key: ex.features.feature[sample_rate_key].int64_list.value.append(8000) old_k = 'oldkey' def _feature_fn(x, s): return tf.expand_dims(tf_frontend.compute_frontend_features( x, s, overlap_seconds=79), axis=-1).numpy().astype(np.float32) do_fn = audio_to_embeddings_beam_utils.ComputeEmbeddingMapFn( name='module_name', module='file.tflite', output_key=0, audio_key=audio_key, sample_rate_key=sample_rate_key, sample_rate=sample_rate, average_over_time=average_over_time, feature_fn=_feature_fn) do_fn.setup() new_k, new_v = next(do_fn.process((old_k, ex))) self.assertEqual(new_k, old_k) expected_shape = (1, BASE_SHAPE_[1]) if average_over_time else BASE_SHAPE_ self.assertEqual(new_v.shape, expected_shape)
def test_compute_embedding_map_fn(self, average_over_time, sample_rate_key, sample_rate): # Establish required key names. audio_key = 'audio_key' # Construct the tf.train.Example test data. ex = tf.train.Example() ex.features.feature[audio_key].float_list.value.extend( np.zeros(2000, np.float32)) if sample_rate_key: ex.features.feature[sample_rate_key].int64_list.value.append(8000) old_k = 'oldkey' do_fn = audio_to_embeddings_beam_utils.ComputeEmbeddingMapFn( name='module_name', module='@loc', output_key='unnecessary', audio_key=audio_key, sample_rate_key=sample_rate_key, sample_rate=sample_rate, average_over_time=average_over_time) do_fn.setup() new_k, new_v = next(do_fn.process((old_k, ex))) self.assertEqual(new_k, old_k) expected_shape = (1, BASE_SHAPE_[1]) if average_over_time else BASE_SHAPE_ self.assertEqual(new_v.shape, expected_shape)
def test_compute_embedding_map_fn_custom_call(self, average_over_time, sample_rate_key, sample_rate): # Establish required key names. audio_key = 'audio_key' custom_call_shape = (5, 25) # Custom call function for embedding generation. def test_call_fn(audio_samples, sample_rate, module_location, output_key, name): """Mock waveform-to-embedding computation.""" del audio_samples, sample_rate, module_location, output_key, name return np.zeros(custom_call_shape, dtype=np.float32) # Construct the tf.train.Example test data. ex = tf.train.Example() ex.features.feature[audio_key].float_list.value.extend( np.zeros(2000, np.float32)) if sample_rate_key: ex.features.feature[sample_rate_key].int64_list.value.append(8000) old_k = 'oldkey' do_fn = audio_to_embeddings_beam_utils.ComputeEmbeddingMapFn( name='module_name', module='@loc', output_key='unnecessary', audio_key=audio_key, sample_rate_key=sample_rate_key, sample_rate=sample_rate, average_over_time=average_over_time, module_call_fn=test_call_fn, setup_fn=lambda _: None) do_fn.setup() new_k, new_v = next(do_fn.process((old_k, ex))) self.assertEqual(new_k, old_k) expected_shape = ( 1, custom_call_shape[1]) if average_over_time else custom_call_shape self.assertEqual(new_v.shape, expected_shape)