Repository: tensorflow/addons

Issue:
reuse the instance of class BeamSearchDecoder  more easily
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
- TensorFlow version and how it was installed (source or binary):
- TensorFlow-Addons version and how it was installed (source or binary):
- Python version: 3.6
- Is GPU used? (yes):

**Describe the bug**

A clear and concise description of what the bug is.

**Code to reproduce the issue**

Provide a reproducible test case that is the bare minimum necessary to generate the problem.

**Other info / logs**

I'm new for TF. I trained a seq2seq model based on LSTM for NLG, then to predict with Seq2SeqPredict, code as follows. The decoder of Seq2SeqPredict is the instance of `BeamSearchDecoder`, and My Error raised in it.
```
class Seq2SeqPredict(tf.keras.Model):
    def __init__(self, Ty, beam_width, start_token, end_token, seq2seq_train):
        super(Seq2SeqPredict, self).__init__()
        
        ## encoder
        self.encoder_embedding = seq2seq_train.encoder_embedding
        self.encoder_rnn = seq2seq_train.encoder_rnn
        
        ## decoder embedding
        self.decoder_embedding_mtx = seq2seq_train.decoder_embedding.variables[0]
        ## Output
        # # self.decoder_dense_layer = seq2seq_train.decoder_dense_layer
        ## attention_mechanism
        self.decoder_attention_mechanism = seq2seq_train.decoder_attention_mechanism
        
        ## rnn_cell
        self.decoder_rnn_cell = seq2seq_train.decoder_rnn_cell
        
        ## beam_search_decoder
        self.beam_search_decoder = tfa.seq2seq.BeamSearchDecoder(
            cell=self.decoder_rnn_cell,
            beam_width=beam_width,
            output_layer=seq2seq_train.decoder_dense_layer,
        )
        
        ## other attributions
        self.beam_width = beam_width
        self.maximum_iterations = Ty * 2
        self.start_token = start_token
        self.end_token = end_token
        
    def _decoder_initial_state(self, size, state):
        _initial_state = self.decoder_rnn_cell.get_initial_state(batch_size=size*self.beam_width, dtype=tf.float32)
        decoder_initial_state = _initial_state.clone(cell_state=state)
        return decoder_initial_state
    
    def _generate_lines(self, start_tokens, tiled_state):
        beam_width = self.beam_width
        (first_finished, first_inputs, first_state) = self.beam_search_decoder.initialize(self.decoder_embedding_mtx,
                                                                                         start_tokens=start_tokens,
                                                                                         end_token=self.end_token,
                                                                                         initial_state=self._decoder_initial_state(1, state=tiled_state),
                                                                                         )
        inputs = first_inputs
        state = first_state
        pred_mtx, path_mtx = [], []

        for j in range(self.maximum_iterations):
            outputs, next_state, next_inputs, finished = self.beam_search_decoder.step(j, inputs, state)
            inputs = next_inputs
            state = next_state

            pred_mtx.append(outputs.predicted_ids)
            path_mtx.append(outputs.parent_ids)
        
        return pred_mtx, path_mtx

    def call(self, inputs, training=False):
        encoder_input = inputs
        # encoder
        encoder_x = self.encoder_embedding(encoder_input)
        a, a_tx, c_tx = self.encoder_rnn(encoder_x)
        
        beam_width = self.beam_width
        # beam_search_decoder
        ## tiled
        tiled_a = tfa.seq2seq.tile_batch(a, multiplier=self.beam_width)
        tiled_a_tx = tfa.seq2seq.tile_batch(a_tx, multiplier=self.beam_width)
        tiled_c_tx = tfa.seq2seq.tile_batch(c_tx, multiplier=self.beam_width) 
        start_tokens = tf.fill([1], self.start_token)

        self.decoder_attention_mechanism.setup_memory(tiled_a)

        pred_mtx, path_mtx = self._generate_lines(start_tokens, [tiled_a_tx, tiled_c_tx])
        return pred_mtx, path_mtx
```
when I predict sequences twicely with the same instance of class Seq2SeqPredict, an error raised "embedding and embedding_fn cannot be provided at same time". 

Reason is that instance would create automatically a `_embedding_fn` with `embedding` in the first prediction without any notice.

I think some action should be taken, as follows:
1.  One set `embedding` in function `initialize()`, so he want to replace old `_embedding_fn` with it. So replace older.

2. Given performance, it's better to set `embedding` when creating new instance. So append a message to the note.

Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.



Discussion:
Tensorflow 2.1.0    installed by Anaconda
Tensorflow_addons 0.9.1   installed by Anaconda
Python 3.6.9    install by Anaconda
GPU      yes!

You can set `embedding_fn` in the constructor of `BeamSearchDecoder`. Would that work for you?

A typical embedding function would look like this:

```python
lambda ids: tf.nn.embedding_lookup(embedding, ids)
```
yes, I had done as you said before creating this issue.&nbsp; But it's better giving advice in note or other.&nbsp;



---Original---
From: "Guillaume Klein"<notifications@github.com&gt;
Date: Fri, May 22, 2020 16:11 PM
To: "tensorflow/addons"<addons@noreply.github.com&gt;;
Cc: "Author"<author@noreply.github.com&gt;;"MyTool"<1607350868@qq.com&gt;;
Subject: Re: [tensorflow/addons] reuse the instance of class BeamSearchDecoder  more easily (#1868)




 
You can set embedding_fn in the constructor of BeamSearchDecoder. Would that work for you?
 
A typical embedding function would look like this:
 lambda ids: tf.nn.embedding_lookup(embedding, ids)
 
—
You are receiving this because you authored the thread.
Reply to this email directly, view it on GitHub, or unsubscribe.
@qlzh727 What is the rationale for having the `embedding` argument in the call method and not only in the constructor?
If we set `embedding` in `initialize` method, the instance of BeamSearchDecoder run only once. It is one-off.

Resolving patch:
diff --git a/tensorflow_addons/seq2seq/beam_search_decoder.py b/tensorflow_addons/seq2seq/beam_search_decoder.py
--- a/tensorflow_addons/seq2seq/beam_search_decoder.py
+++ b/tensorflow_addons/seq2seq/beam_search_decoder.py
@@ -663,8 +663,9 @@ def __init__(
           cell: A layer that implements the `tf.keras.layers.AbstractRNNCell`
             interface.
           beam_width:  Python integer, the number of beams.
-          embedding_fn: A callable that takes a vector tensor of `ids`
-            (argmax ids).
+          embedding_fn: A callable that takes a `int32` `Tensor` of token IDs
+            and returns embedding tensors. If set, the `embedding` argument in
+            the decoder call should be set to `None`.
           output_layer: (Optional) An instance of `tf.keras.layers.Layer`,
             i.e., `tf.keras.layers.Dense`.  Optional layer to apply to the RNN
             output prior to storing the result or sampling.
@@ -696,24 +697,32 @@ def initialize(self, embedding, start_tokens, end_token, initial_state):
         """Initialize the decoder.
 
         Args:
-          embedding: A tensor from the embedding layer output, which is the
-            `params` argument for `embedding_lookup`.
-          start_tokens: `int32` vector shaped `[batch_size]`, the start tokens.
-          end_token: `int32` scalar, the token that marks end of decoding.
-          initial_state: A (possibly nested tuple of...) tensors and
-          TensorArrays.
+          embedding: A `Tensor` (or `Variable`) to pass as the `params` argument
+            for `tf.nn.embedding_lookup`. This overrides `embedding_fn` set in
+            the constructor.
+          start_tokens: Start the decoding from these tokens.
+            A `int32` `Tensor` of shape `[batch_size]`.
+          end_token: The token that marks the end of decoding.
+            A `int32` scalar `Tensor`.
+          initial_state: The initial cell state as a (possibly nested) structure
+            of `Tensor` and `TensorArray`.
+
         Returns:
           `(finished, start_inputs, initial_state)`.
+
         Raises:
+          ValueError: If `embedding` is `None` and `embedding_fn` was not set
+            in the constructor.
           ValueError: If `start_tokens` is not a vector or `end_token` is not a
             scalar.
         """
-        if embedding is not None and self._embedding_fn is not None:
+        if embedding is not None:
+            self._embedding_fn = lambda ids: tf.nn.embedding_lookup(embedding, ids)
+        elif self._embedding_fn is None:
             raise ValueError(
-                "embedding and embedding_fn cannot be provided at same time"
+                "You should either pass an embedding variable when calling the "
+                "BeamSearchDecoder or set embedding_fn in the constructor."
             )
-        elif embedding is not None:
-            self._embedding_fn = lambda ids: tf.nn.embedding_lookup(embedding, ids)
 
         self._start_tokens = tf.convert_to_tensor(
             start_tokens, dtype=tf.int32, name="start_tokens"


Test patch:
diff --git a/tensorflow_addons/seq2seq/tests/beam_search_decoder_test.py b/tensorflow_addons/seq2seq/tests/beam_search_decoder_test.py
--- a/tensorflow_addons/seq2seq/tests/beam_search_decoder_test.py
+++ b/tensorflow_addons/seq2seq/tests/beam_search_decoder_test.py
@@ -280,6 +280,20 @@ def test_eos_masking():
         np.testing.assert_allclose(masked[1][2][i], np.finfo("float32").min)
 
 
+def test_missing_embedding_fn():
+    batch_size = 6
+    beam_width = 4
+    cell = tf.keras.layers.LSTMCell(5)
+    decoder = beam_search_decoder.BeamSearchDecoder(cell, beam_width=beam_width)
+    initial_state = cell.get_initial_state(
+        batch_size=batch_size * beam_width, dtype=tf.float32
+    )
+    start_tokens = tf.ones([batch_size], dtype=tf.int32)
+    end_token = tf.constant(2, dtype=tf.int32)
+    with pytest.raises(ValueError):
+        decoder(None, start_tokens, end_token, initial_state)
+
+
 def test_beam_step():
     batch_size = 2
     beam_width = 3