5

I am using the python huggingface transformers library for a text-generation model. I need to know how to implement the stopping_criteria parameter in the generator() function I am using.

I found the stopping_criteria parameter in this documentation: https://huggingface.co/transformers/main_classes/pipelines.html#transformers.TextGenerationPipeline

The problem is, I just dont know how to implement it.

My Code:

from transformers import pipeline
generator = pipeline('text-generation', model='EleutherAI/gpt-neo-125M')
stl = StoppingCriteria(['###'])
res = generator(prompt, do_sample=True,stopping_criteria = stl)
Developer
  • 98
  • 1
  • 9
  • 2
    I was also curious about this. I'd like to be able to provide a particular stopping token (other than the EOS token). Did you work this out? – jbm Feb 06 '22 at 16:35
  • How do we resolve "StoppingCriteria() takes no arguments"? It says it's an abstract class; so do we need to define a new class and if so which methods do we override? – pete Nov 21 '22 at 00:11

1 Answers1

1

These two approaches worked for me. your_condition is True when you want to stop.

class CustomStoppingCriteria(StoppingCriteria):
    def __init__(self):
        pass
    
    def __call__(self, input_ids: torch.LongTensor, score: torch.FloatTensor, **kwargs) -> bool:
        return your_condition

stopping_criteria = StoppingCriteriaList([CustomStoppingCriteria()])

OR

def custom_stopping_criteria(input_ids: torch.LongTensor, score: torch.FloatTensor, **kwargs) -> bool:
    return your_condition

stopping_criteria = StoppingCriteriaList([custom_stopping_criteria])
vv22
  • 11
  • 2