Spaces:
Running
Running
from transformers import StoppingCriteria | |
# Define a stopping condition for generation | |
class SpecificStringStoppingCriteria(StoppingCriteria): | |
def __init__(self, tokenizer, stop_strings, input_len): | |
self.tokenizer = tokenizer | |
self.stop_strings = stop_strings | |
self.input_len = input_len | |
def __call__(self, input_ids, scores, **kwargs): | |
current_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)[self.input_len:] | |
return any(stop_string in current_text for stop_string in self.stop_strings) |