def get_naive_policy(self, budget=250):
assert budget <= self.num_steps
# We use budget+1 because a linspace contains the last step.
naive_policy = ardm_utils.integer_linspace(0, self.num_steps, budget+1)
# Last index does not need to be in policy.
naive_policy = naive_policy[:-1]
return naive_policy
def print_language_sample_chain(chain, decode_tokens, keep_frames=5):
"""Prints a language chain to visualize the generation process."""
if len(chain) > keep_frames:
linspace = ardm_utils.integer_linspace(0, len(chain), keep_frames)
chain = chain[linspace]
exemplars = [decode_tokens(batch[0]) for batch in chain]
logging.info('Text sampling chain:')
logging.info('\n'.join(exemplars))
def get_naive_policy(self, budget_per_stage=500):
assert budget_per_stage <= self.num_steps_per_stage
# We use budget_per_stage+1 because a linspace contains the last step.
naive_policy = ardm_utils.integer_linspace(0, self.num_steps_per_stage,
budget_per_stage + 1)
# Last index does not need to be in policy
naive_policy = naive_policy[:-1]
naive_policy = jnp.repeat(naive_policy[None, :],
self.num_stages,
axis=0)
return naive_policy