def add()

in lm_human_preferences/utils/core.py [0:0]

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

    def add(self, **data):
        """Add new data to the end of the buffer, dropping old data if we exceed capacity."""
        # Check input shapes
        if data.keys() != self._vars.keys():
            raise ValueError('data.keys() = %s != %s' % (sorted(data.keys()), sorted(self._vars.keys())))
        first = next(iter(data.values()))
        pre = first.shape[:1]
        for k, d in data.items():
            try:
                d.shape.assert_is_compatible_with(pre.concatenate(self._vars[k].shape[1:]))
            except ValueError as e:
                raise ValueError('%s, key %s' % (e, k))
        # Enqueue
        n = tf.shape(first)[0]
        capacity = self._capacity
        i0 = (self._total.assign_add(n) - n) % capacity
        i0n = i0 + n
        i1 = tf.minimum(i0n, capacity)
        i2 = i1 % capacity
        i3 = i0n % capacity
        slices = slice(i0, i1), slice(i2, i3)
        sizes = tf.stack([i1 - i0, i3 - i2])
        assigns = [self._vars[k][s].assign(part)
                   for k,d in data.items()
                   for s, part in zip(slices, tf.split(d, sizes))]
        return tf.group(assigns)