Implement TopP, TopK and Beam samplers

[chenmoneygithub]

A few things covered:

• Implement TopP, TopK and Beam samplers
• Add "Sampler" suffix to our sampler class names.
• Add run_eagerly option to our sampler to align with model.compile().

One special thing to note:

• In Beam sampler, there is a strange error on batch_size=1 case, which is after the first iteration, the shape of beam_probs changes from [1, None] to [None, None], so we add shape_invariants specifically for beam sampler.

[jbischof]

Thanks @chenmoneygithub looking great! The biggest thing is unifying our docstring style

[jbischof]

@chenmoneygithub I really don't think these arg blocks match the rest of the library's style. I don't think there's a perfect answer but I'd prefer it not to be obvious who wrote what 🥗

[chenmoneygithub]

Thanks! I checked our code base, it appears we have a few code having this style:

• maskedLmhead: link
• PositionEmbedding: link
• SinePositionEmbedding: link
• Generation utils: link

I feel like giving these numbers an explicit meaning makes the code longer but kinda easier to understand.

[jbischof]

It's a totally valid opinion @chenmoneygithub but it's also important for the code to have a unified style rather than each contributor producing different looking code. If there's an example that's unclear without named params I'm open to trying something different but otherwise I'm hoping we can compromise!

[mattdangerw]

Taking steps towards a more unified style (and then reflecting that in our style guide) sgtm. What are the main places this differs, besides the constants at the top?

[jbischof]

Basically what I asked for in #658 was the current thought. A bit less script-like and a little more "drop this line in colab and see what we're talking about".

[mattdangerw]

I could see us replacing the "model" with something like tf.random.uniform(shape, minval=-1, maxval=1). It is kind of weird to me that we show a whole model, that is trainable but randomly initialized (so results will be random anyway), and not even sequence aware so would never really perform even if your trained it. For a new user this seems a bit of a red herring.

Would be more concise to do something like:

    def token_probability_fn(inputs, mask):
        return tf.random.uniform(...) # Replace with a real model!

[chenmoneygithub]

I would like to keep the model part so that the example is closer to real use cases.

I will unify the docstring to move those hypers inline.

[mattdangerw]

To me the model falls into an "uncanny valley" of code examples. It's not something that will actually work, yet also not clearly a random dummy data. As a newbie I worry I would not understand, first, that results will be random, and second, that this is a "bad model" for the task.

Fine to merge as is, but I hope we can play around with some improvements down the road.

[jbischof]

Do we need two flags or just one then? What happens if they are both False---what is a "non-XLA" graph? I'm confused!

[chenmoneygithub]

This is following the style of model.compile(): code link

non-XLA graph is quite common, anything annotated with tf.function without jit_compile=True is a non-xla graph.

[jbischof]

Yes, agree, was just trying to understand the difference

[mattdangerw]

Thanks Chen! Left some comments and a few questions.

[mattdangerw]

return immediately

[mattdangerw]

The main change I could see us making if we want to, is splitting the overload points into sample and sample_step (I think some discussion of this on other PR?).

Overload sample if you want to control the whole process (e.g. beam search). Overload sample_step if you want to control simply going from one probability distribution to one sample. E.g. the body of this class could look like:

def sample_step(preds):
     sorted_preds, sorted_indices = tf.math.top_k(
          pred, k=tf.shape(preds)[1], sorted=True
      )
      # Calculate cumulative probability distribution.
      cumulative_probs = tf.math.cumsum(sorted_preds, axis=-1)
      # Create a mask for the tokens to keep.
      keep_mask = cumulative_probs <= self.p
      # Shift to include the last token that exceed p.
      shifted_keep_mask = tf.concat(
          [tf.ones_like(keep_mask[:, :1]), keep_mask[:, :-1]], axis=-1
      )
      # Filter out unmasked tokens and sample from filtered distribution.
      probs = tf.where(
          shifted_keep_mask,
          sorted_preds,
          tf.zeros(tf.shape(pred), dtype=sorted_preds.dtype),
      )
      sorted_next_token = tf.random.categorical(
          tf.math.log(probs), 1, seed=self.seed
      )
      return tf.gather_nd(
          sorted_indices, sorted_next_token, batch_dims=1
      )

This might improve the readability of our simple samples, while still keeping full extensibility.

[chenmoneygithub]

As long as we can make BeamSampler an outlier, I am down to this refactoring!

Actually we can move one step further to only leave get_next_token open, which takes in a prob distribution and returns the next token, and the rest updating logic can be shared across those samplers.

Reflected this change in the PR.

[jbischof]

Yes, this is what I was hoping for! 🚀

[mattdangerw]

I could see us replacing the "model" with something like tf.random.uniform(shape, minval=-1, maxval=1). It is kind of weird to me that we show a whole model, that is trainable but randomly initialized (so results will be random anyway), and not even sequence aware so would never really perform even if your trained it. For a new user this seems a bit of a red herring.

Would be more concise to do something like:

    def token_probability_fn(inputs, mask):
        return tf.random.uniform(...) # Replace with a real model!

[jbischof]

Just some nits. Looking good!

[jbischof]

Let's inline this arg as well. Or am I missing something?

[chenmoneygithub]

This one is used twice - embedding and dense, and the arg name does not suggest "vocab_size", so I am keeping this one for clarify.

[jbischof]

Do you think sliding window can fit into this paradigm?

[chenmoneygithub]

yea, it should work

[jbischof]

For clarity you could use named args:

print(sampler(prompt, token_probability_fn, max_length=10))

This matches or perhaps enhances the readability of using globals.

[chenmoneygithub]

good call, done

[mattdangerw]

Left a few more comments re get_config. +approval, as I am fine to merge after changes to get something landed. But hope we will stay open to changes here as we dig more into the generative case!

[jbischof]

Wait for @fchollet, he wants to take a look and can make the final call on class names.

[jbischof]

Short class names seem to fit in the Keras ecosystem:

sampler=keras_nlp.samplers.Beam(num_beams=3),
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=keras.metrics.SparseCategoricalAccuracy(),
kernel_initializer=initializers.RandomNormal(stddev=0.01),
activation=activations.relu,

[chenmoneygithub]

Let's merge this one to unblock the TFLite sprint, I will sync with Francois offline regarding the string identifier.

[fchollet]

This should raise an error I suppose?