Method batch_decode

opencompass/models/interntrain.py:480–500

(self,
                     outputs,
                     eos_token_ids: List[int],
                     stopping_criteria: List[str] = [])

        return torch.LongTensor(tokens).cuda()

    def batch_decode(self,
                     outputs,
                     eos_token_ids: List[int],
                     stopping_criteria: List[str] = []):
        # outputs: bsz, seq_len
        output_text = []
        outputs = outputs.tolist()
        for output in outputs:
            # cut off by eos_token_ids
            eos_idx = len(output)
            for eos_id in eos_token_ids:
                if eos_id in output:
                    eos_idx = min(output.index(eos_id), eos_idx)
            text = self.tokenizer.decode(output[:eos_idx])
            if self.end_str is not None:
                text = text.split(self.end_str)[0]
            for stop_word in stopping_criteria:
                text = text.split(stop_word)[0]
            output_text.append(text)

        return output_text