def download_labels()

in lm_human_preferences/train_reward.py [0:0]

• 17 lines of code
• 6 McCabe index (conditional complexity)

def download_labels(source, label_type, question_schemas, total_labels, comm):
    schemas = {**question_schemas, **label_type.label_schemas()}

    """
    if self.is_root:
        with tf.device('cpu:0'):
            self._enqueue_phs = {
                name: tf.placeholder(name=name, dtype=schema.dtype, shape=(None,) + schema.shape)
                for name, schema in self.schemas.items()
            }
            self._enqueue_answers = self.answer_queue.enqueue_many(self._enqueue_phs)
    else:
        self._enqueue_phs = None
        self._enqueue_answers = None
    """

    # TODO: download on just one rank?  then do: labels = utils.mpi_bcast_tensor_dict(labels, comm=comm)
    if source != 'test':
        with open(gcs.download_file_cached(source, comm=comm)) as f:
            results = json.load(f)
            print('Num labels found in source:', len(results))
    else:
        results = [
            {
                name: np.zeros(schema.shape, dtype=schema.dtype.as_numpy_dtype)
                for name, schema in schemas.items()
            }
            for _ in range(50)
        ]

    assert len(results) >= total_labels
    results = results[:total_labels]
    return {k: [a[k] for a in results] for k in schemas.keys()}