def make_query_sampler()

in lm_human_preferences/lm_tasks.py [0:0]

• 21 lines of code
• 3 McCabe index (conditional complexity)

def make_query_sampler(*, hparams: TaskHParams, encoder, batch_size: int, mode='train', comm=None):
    if hparams.start_text:
        start_token, = encoder.encode(hparams.start_text)
    else:
        start_token = None

    if hparams.end_text:
        end_token, = encoder.encode(hparams.end_text)
    else:
        end_token = None

    data = datasets.get_dataset(hparams.query_dataset).tf_dataset(
        sequence_length=hparams.query_length, mode=mode, comm=comm, encoder=encoder,
        start_token=start_token, end_token=end_token,
    )
    data = data.map(lambda d: tf.cast(d['tokens'], tf.int32))
    data = data.batch(batch_size, drop_remainder=True)

    context_iterator = data.make_one_shot_iterator()

    def sampler(scope=None):
        with tf.name_scope(scope, 'sample_corpus'):
            context_tokens = context_iterator.get_next()
            return dict(tokens=context_tokens)
    return sampler