Method get_flat_data

train/dataloader.py:634–656

Return a flat list of examples given a list of prompts that index self.full_data. Assumes that there are a list of scores.

(self, prompts)

class ScoreDataLoader(UnpairedPreferenceDataLoader):
    def get_flat_data(self, prompts):
        """
        Return a flat list of examples given a list of prompts that index self.full_data.
        Assumes that there are a list of scores.
        """
        flat_data = []
        prev_status = 'rejected'

        for prompt in prompts:
            example = self.full_data[prompt]

            if self.max_prompt_count:
                example.pairs = self.rng.sample(example.pairs, min(self.max_prompt_count, len(example.pairs)))

            # for oasst, lower scores are better, so rank 0 is the best response and rank n is the worst
            if prev_status == 'rejected':
                flat_data.append((example, example.generations[np.argmin(example.scores)], 'chosen'))
            else:
                flat_data.append((example, example.generations[np.argmax(example.scores)], 'rejected'))

            prev_status = flat_data[-1][-1]

        return flat_data


class HalfPrefDataLoader(UnpairedPreferenceDataLoader):