def get_orderings()

in rbr_weight_fitter.py [0:0]

• 18 lines of code
• 7 McCabe index (conditional complexity)

    def get_orderings(self):
        def apply_ordering(i1, i2):
            ex1, ex2 = self.examples[i1], self.examples[i2]
            if not ex1.convo_prompt == ex2.convo_prompt:
                return None
            if ex2.completion_label in self.orderings[ex1.completion_label]:
                margin = self.completion_label_margin.get(ex1.completion_label, 1)
                return [i1, i2, margin]
            elif ex1.completion_label in self.orderings[ex2.completion_label]:
                margin = self.completion_label_margin.get(ex2.completion_label, 1)
                return [i2, i1, margin]
            return None

        pairs = list(combinations(range(len(self.examples)), 2))
        pairs = [apply_ordering(i1, i2) for i1, i2 in pairs]
        pairs = [p for p in pairs if p is not None]
        np.random.default_rng().shuffle(pairs)
        pairs = torch.tensor(pairs, dtype=torch.int64)
        return pairs