我正在为
transformers
模型使用 python huggingface text-generation
库。我需要知道如何在我正在使用的 stopping_criteria
函数中实现 generator()
参数。
我在本文档中找到了
stopping_criteria
参数:
https://huggingface.co/transformers/main_classes/pipelines.html#transformers.TextGenerationPipeline
问题是,我只是不知道如何实现它。
我的代码:
from transformers import pipeline
generator = pipeline('text-generation', model='EleutherAI/gpt-neo-125M')
stl = StoppingCriteria(['###'])
res = generator(prompt, do_sample=True,stopping_criteria = stl)
这两种方法对我有用。
your_condition
是 True 当你想停下来的时候。
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()])
或
def custom_stopping_criteria(input_ids: torch.LongTensor, score: torch.FloatTensor, **kwargs) -> bool:
return your_condition
stopping_criteria = StoppingCriteriaList([custom_stopping_criteria])