import torch
from transformers import StoppingCriteria


class StoppingCriteriaSub(StoppingCriteria):

    def __init__(self, stops=[]):
        StoppingCriteria.__init__(self),

    def __call__(
        self,
        input_ids: torch.LongTensor,
        scores: torch.FloatTensor,
        stops=[]
    ):
        self.stops = stops
        for i in range(len(stops)):
            self.stops = self.stops[i]
