Class CodeBlockStopper

fastchat/model/model_codet5p.py:50–58

    )

    class CodeBlockStopper(StoppingCriteria):
        def __call__(
            self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
        ) -> bool:
            # Code-completion is open-end generation.
            # We check \n\n to stop at end of a code block.
            if list(input_ids[0][-2:]) == [628, 198]:
                return True
            return False

    gen_kwargs = dict(
        **encoding,