How to implement `stopping_criteria` parameter in transformers library?

Viewed 293

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)
1 Answers

According to documentation, to use stopping_criteria you need to pass token ID instead of string.

How to add your custom token (like '###') described here

Related