llama大模型提前停止策略，实现工具调用——以Llama3为例

在大模型的generate过程中为了实现工具调用功能，模型输出到了指定的token需要执行工具的时候，需要模型的generate停止。model.generate()的参考链接：https://github.com/huggingface/transformers/blob/v4.20.1/src/transformers/generation_utils.py#L844

其中有一个参数：stopping_criteria，用于判断模型输出是否应该停止，下面介绍具体方案。

Step1：重新定义一个类继承StoppingCriteria并定义停止逻辑

class KeyWordOne_StoppingCriteria(StoppingCriteria):
    def __init__(self,keyword,tokenizer,device):
        self.keyword = tokenizer.encode(keyword,add_special_tokens = False,return_tensors = 'pt').squeeze().to(device)
    def __call__(self,input_ids,scores,**kwards):
        if len(input_ids[0]) < len(self.keyword):
            return False
        if input_ids[0][len(input_ids[0] - len(self.keyword):].equal(self.keyword):
            return True
        else:
            return False

上面的代码中，构造函数一般用来确定我们要停止的tokens
call调用函数是用来确定停止逻辑的。

Step2：实例化类，并将类传入到generate函数中即可实现提前停止。

from self import KeyWordOne_StoppingCriteri
from transformers import StoppingCriteriaList
stopcrieria = KeyWordOne_StoppingCriteri("<Retrieval>",tokenizer=tokenizer,device=device)
data = [{"content": "针对下列问题，请结合你的知识判断是否需要检索外部数据库以协助回答。如果需要检索，请你给出<Retrieval>标签代表需要检索外部数据库，如果你给出了<Retrieval>标签，那么在它和<Retrieval/>标签中间代表的是从外部数据库检索到的数据，不需要你进行预测。请你根据给出的知识或者自身的能力回答问题。请注意：1、有些问题可能无需外部资料即可回答。2、所提供的资料可能包含与问题无关的信息，请忽略并基于你已有的知识回答问题。", "role": "system"}, 
        {"content": "请在法律条文中上仔细的告知我，如果我在澳门特别行政区设立一家公司，会受到中国大陆财政政策的影响吗，还是澳门有自己独立的财政政策？", "role": "user"}]
data = tokenizer.apply_chat_template(data,add_generation_prompt=True,tokenize=False)
generate_ids = tokenizer(data,return_tensors='pt').to(device)
stop_ids = tokenizer.encode("<Retrieval>",return_tensors='pt')
generated_ids = model.generate(**generate_ids,stopping_criteria=StoppingCriteriaList([stopcrieria]))
tokenizer.batch_decode(generated_ids,skip_special_tokens=True)[0]

上面的代码中，使用StoppingCriteriaList方法包裹了我们实例化的自定义的StoppingCriteria的子类，并用列表包裹起来。

缺点：感觉这样的方案使得模型的推理失去了并行性。

总结