Falcon 7B与LangChain：构建具备对话记忆的智能聊天机器人

一、前言

本文我们将介绍使用 Falcon 7B 和 LangChain 来构建一个保留对话记忆的聊天机器人。通过利用单个 T4 GPU 并以8bit（约6个token/秒）的速度加载模型，以达到不错的模型性能效果。

首先，我们将介绍模型停止准则。通过检测LLM开始"胡言乱语"的时候，并停止生成，我们可以避免生成无意义或混乱的回复。这样可以提高聊天机器人的可读性和可理解性。
其次，我们将讨论清理输出的方法。有时候，LLMs会输出奇怪或多余的标记，这可能会影响到对话的连贯性和准确性。我将向您展示如何清除这些标记，使得输出更加干净和准确。
最后，我们将介绍存储聊天记录的重要性。通过使用内存来保存对话历史，我们可以确保聊天机器人记住之前的对话内容，从而更好地理解用户的需求和上下文。这将有助于提供更加个性化和连贯的回复。

我们将使用Jupyter Notebook来进行代码演示。

二、设置

安装所需的依赖：

!pip install -Uqqq pip --progress-bar off
!pip install -qqq bitsandbytes==0.40.0 --progress-bar off
!pip install -qqq torch==2.0.1 --progress-bar off
!pip install -qqq git+https://github.com/huggingface/transformers --progress-bar off
!pip install -qqq accelerate==0.21.0 --progress-bar off
!pip install -qqq xformers==0.0.20 --progress-bar off
!pip install -qqq einops==0.6.1 --progress-bar off
!pip install -qqq langchain==0.0.233 --progress-bar off

导入必要的库:

注意：由于这里使用的Google Colab环境，我们需要将原始权重转换为更小的块，以便使用Accelerate高效加载。由于原始权重很大，在环境中消耗了所有的RAM，导致免费的GPU资源无法直接运行tiiuae/falcon-7b-instruct模型，会出现内存占满崩溃的问题。

幸运的是，我们可以利用量化技术对深度学习模型的量化和加速进行自定义配置。HuggingFace与BitsAndBytes库进行了集成，以便更好地支持深度学习模型的量化和加速。

import re
import warnings
from typing import List

import torch
from langchain import PromptTemplate
from langchain.chains import ConversationChain
from langchain.chains.conversation.memory import ConversationBufferWindowMemory
from langchain.llms import HuggingFacePipeline
from langchain.schema import BaseOutputParser
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    StoppingCriteria,
    StoppingCriteriaList,
    pipeline,
)

warnings.filterwarnings("ignore", category=UserWarning)

三、加载模型

接下来，我们可以直接从 Hugging Face 模型库中加载模型，这里我们选择重新分片的版本来加载模型vilsonrodrigues/falcon-7b-instruct-sharded，可以直接适用于比较低RAM环境（例如Colab、Kaggle）中的safetensors。：

MODEL_NAME = "vilsonrodrigues/falcon-7b-instruct-sharded" 

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME, trust_remote_code=True, device_map="auto", quantization_config=quantization_config
)
model = model.eval()

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
print(f"Model device: {model.device}")

请注意，我们以4位量化模式加载模型。将会减少内存占用从而加快推理速度。我们还使用device_map参数将模型加载到GPU上。

3.1、配置模型

对于文本生成，我们将使用自定义配置：

generation_config = model.generation_config
generation_config.temperature = 0
generation_config.num_return_sequences = 1
generation_config.max_new_tokens = 256
generation_config.use_cache = False
generation_config.repetition_penalty = 1.7
generation_config.pad_token_id = tokenizer.eos_token_id
generation_config.eos_token_id = tokenizer.eos_token_id
generation_config

GenerationConfig {
    "_from_model_config": true,
    "bos_token_id": 1,
    "eos_token_id": 11,
    "max_new_tokens": 256,
    "pad_token_id": 11,
    "repetition_penalty": 1.7,
    "temperature": 0,
    "transformers_version": "4.30.0",
    "use_cache": false
}

将temperature设置为0，可以获得比较确定的结果。

在文本生成中，temperature参数控制着生成文本的多样性。当temperature接近0时，生成的结果更加确定和一致。因此，将temperature设置为0可以确保生成的结果具有确定性。

将repetition_penalty设置为1.7，以减少模型重复自身的机会（但不完全消除）。

repetition_penalty参数用于控制模型生成文本时是否倾向于避免重复使用相同的词语或短语。通过设置repetition_penalty参数，可以降低模型生成重复内容的概率，但并不能完全消除重复出现的可能性。

打印模型的配置信息：

FalconConfig {
  "_name_or_path": "vilsonrodrigues/falcon-7b-instruct-sharded",
  "alibi": false,
  "apply_residual_connection_post_layernorm": false,
  "architectures": [
    "FalconForCausalLM"
  ],
  "attention_dropout": 0.0,
  "auto_map": {
    "AutoConfig": "vilsonrodrigues/falcon-7b-instruct-sharded--configuration_falcon.FalconConfig",
    "AutoModel": "vilsonrodrigues/falcon-7b-instruct-sharded--modeling_falcon.FalconModel",
    "AutoModelForCausalLM": "vilsonrodrigues/falcon-7b-instruct-sharded--modeling_falcon.FalconForCausalLM",
    "AutoModelForQuestionAnswering": "vilsonrodrigues/falcon-7b-instruct-sharded--modeling_falcon.FalconForQuestionAnswering",
    "AutoModelForSequenceClassification": "vilsonrodrigues/falcon-7b-instruct-sharded--modeling_falcon.FalconForSequenceClassification",
    "AutoModelForTokenClassification": "vilsonrodrigues/falcon-7b-instruct-sharded--modeling_falcon.FalconForTokenClassification"
  },
  "bias": false,
  "bos_token_id": 11,
  "eos_token_id": 11,
  "hidden_dropout": 0.0,
  "hidden_size": 4544,
  "initializer_range": 0.02,
  "layer_norm_epsilon": 1e-05,
  "model_type": "falcon",
  "multi_query": true,
  "new_decoder_architecture": false,
  "num_attention_heads": 71,
  "num_hidden_layers": 32,
  "num_kv_heads": 71,
  "parallel_attn": true,
  "quantization_config": {
    "bnb_4bit_compute_dtype": "float16",
    "bnb_4bit_quant_type": "nf4",
    "bnb_4bit_use_double_quant": true,
    "llm_int8_enable_fp32_cpu_offload": false,
    "llm_int8_has_fp16_weight": false,
    "llm_int8_skip_modules": null,
    "llm_int8_threshold": 6.0,
    "load_in_4bit": true,
    "load_in_8bit": false
  },
  "torch_dtype": "bfloat16",
  "transformers_version": "4.30.0",
  "use_cache": true,
  "vocab_size": 65024
}

3.2、使用模型

接下来可以开始使用模型了。首先，定义了一个包含友好对话开头和用户提问的字符串变量"prompt"。然后，使用分词器(tokenizer)对prompt进行编码，转换为模型可理解的输入格式。接下来，将编码后的input_ids传递给模型，并放置在相同设备上以确保推理一致性。最后，使用model.generate函数生成对话回复，传入input_ids作为输入，并指定generation_config参数来配置生成文本的选项。

prompt = """
以下是一个人类和AI之间友好的对话。AI非常健谈，并提供了许多具体的细节。

当前对话：

Human：Dwight K Schrute是谁？
AI：
""".strip()

input_ids = tokenizer(prompt, return_tensors="pt").input_ids
input_ids = input_ids.to(model.device)

with torch.inference_mode():
    outputs = model.generate(
        input_ids=input_ids,
        generation_config=generation_config,
    )

在进行推理之前，我们将编码后的input_ids放到CUDA设备上。我们可以使用分词器将输出解码为可读的格式：

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)

LLM输出

以下是一个人类和AI之间友好的对话。AI非常健谈，并提供了许多具体的细节。

当前对话：

Human：Dwight K Schrute是谁？
AI：德怀特-K-施鲁特（Dwight K Schrute）是电视剧《办公室》中的一个虚构人物。他由 Rainn Wilson 饰演。
User

输出包含完整的提示和生成的回复。效果看起来还不错，下面我们看看如何改进它。

四、阻止LLM胡言乱语

LLM生成的内容经常会偏离主题，生成与问题无关或毫无意义的回复。虽然这是一个持续的研究挑战，但作为LLM在实际应用中的用户，我们可以通过一些方法来解决这个问题。我们将使用一种名为StoppingCriteria的技术来控制输出，防止模型胡言乱语或产生虚构的问题和对话：

class StopGenerationCriteria(StoppingCriteria):
    def __init__(
        self, tokens: List[List[str]], tokenizer: AutoTokenizer, device: torch.device
    ):
        stop_token_ids = [tokenizer.convert_tokens_to_ids(t) for t in tokens]
        self.stop_token_ids = [
            torch.tensor(x, dtype=torch.long, device=device) for x in stop_token_ids
        ]

    def __call__(
        self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
    ) -> bool:
        for stop_ids in self.stop_token_ids:
            if torch.eq(input_ids[0][-len(stop_ids) :], stop_ids).all():
                return True
        return False

init方法使用分词器将标记转换为相应的标记ID，并将其存储为stop_token_ids。 call方法在生成过程中被调用，以输入ID作为输入。它检查输入ID中最后几个标记是否与任何stoptokenids匹配，表示模型开始生成不希望的回复。如果找到匹配项，则返回True，表示应停止生成。否则，返回False以继续生成。

我们将实现一个模型输出停止准则，当LLM生成以Human:或AI:开头的新标记时进行检测。一旦检测到这样的标记，生成过程将停止以防止产生不希望的输出：

stop_tokens = [["Human", ":"], ["AI", ":"]]
stopping_criteria = StoppingCriteriaList(
    [StopGenerationCriteria(stop_tokens, tokenizer, model.device)]
)

我们将创建一个包含停止准则和生成配置的Pipeline。该Pipeline将处理生成过程，并确保应用停止准则以控制输出：

generation_pipeline = pipeline(
    model=model,
    tokenizer=tokenizer,
    return_full_text=True,
    task="text-generation",
    stopping_criteria=stopping_criteria,
    generation_config=generation_config,
)

llm = HuggingFacePipeline(pipeline=generation_pipeline)

使用我们的Pipeline非常简单，只需将提示传递给Pipeline即可：

res = llm(prompt)
print(res)

Pipeline 输出结果：

德怀特-K-施鲁特（Dwight K Schrute）是电视剧《办公室》中的一个虚构人物。他由 Rainn Wilson 饰演。
User

注意生成文本末尾的`User`？稍后我们会处理它。

五、对话链

为了与LLM进行对话，我们将使用 LangChain 中的 ConversationChain：

chain = ConversationChain(llm=llm)
print(chain.prompt.template)

以下是一个人类和AI之间友好的对话。AI非常健谈，并提供了许多具体的上下文细节。如果AI不知道问题的答案，它会诚实地说出自己不知道。

当前对话：
{history}
Human: {input}
AI:

ConversationChain提供了一个默认的 Prompt 提示，适用于一般情况。然而，对于我们特定的用例来说可能不是最合适的。我们可以使用自定义提示，来更好地满足我们的需求。

5.1、自定义提示

假如你现在需要开一家新公司，需要AI给取一个名字、口号和营销宣传材料。作为乔布斯的超级粉丝，你希望将他独特的风格融入到AI生成的内容中。让我们将这两个元素结合起来，利用AI的力量满足你的创造欲望：

template = """
以下是一个人类和AI之间的对话。这个AI的行为方式完全像《办公室》电视剧中的Dwight K Schrute一样。Dwight是一位经验丰富且非常成功的在线营销人员和销售员。他机智、有说服力、直接而实际。Dwight会帮助完成分配给他的每个营销任务。如果Dwight不知道某个问题的答案，他会真诚地说自己不知道。

当前对话：
{history}
Human：{input}
AI:""".strip()

prompt = PromptTemplate(input_variables=["history", "input"], template=template)

这里定义了适合我们要求的Prompt之后，接下来需要确保我们的聊天机器人在回答当前问题时能够记住我们之前的对话上下文：

memory = ConversationBufferWindowMemory(
    memory_key="history", k=6, return_only_outputs=True
)

chain = ConversationChain(llm=llm, memory=memory, prompt=prompt, verbose=True)

首先，我们创建了一个ConversationBufferWindowMemory对象，它有两个参数：

memory_key: 这是用来标识存储在内存中的对话历史记录的键值。在这里，我们将其设置为"history"。
k: 这是指定要保留的最近消息数量的参数。在这里，我们将其设置为6，表示保留最近的6条消息作为对话历史记录。
returnonlyoutputs: 这是一个布尔值，用于指定是否只返回输出而不返回输入。在这里，我们将其设置为True，表示只返回生成的回复而不返回用户的输入。

然后，我们使用LLM模型和上面创建的ConversationBufferWindowMemory对象来创建ConversationChain对象。它有三个参数：

llm：这是一个预训练的语言模型，用于生成回复。
memory：这是一个用于存储对话历史记录的对象。
prompt：这是一个字符串，用于指定初始提示。
verbose：将verbose参数设置为True，以便在生成回复时打印详细的输出信息。

text = "为一个制造家庭汽车并配备大排量V8发动机的汽车制造商想一个名字。这个名字必须是一个单词，并且易于发音。"
res = chain.predict(input=text)
print(res)

详细输出：

以下是一个人类和AI之间的对话。AI的角色设定是苹果公司的创始人乔布斯，他是一位经验丰富、非常成功的在线营销人员和销售员。他机智、有说服力、直接和实际。乔布斯会帮助完成分配给他的每个营销任务。如果乔布斯不知道某个问题的答案，他会诚实地说自己不知道。

当前对话：
Human：思考一个为普通家庭提供高性价比高质量的智能机器人制造商的名称。这个名字必须是一个单词，并且易于发音。
AI:

Chain输出：

V8
User

结果生成的还可以，只是生成文本末尾多了一个User。我们将在后面来修复这个问题。

5.2、清理输出

为了确保我们的聊天机器人输出干净的内容，我们将通过扩展 LangChain 中的基本OutputParser类来自定义行为。虽然输出解析器通常用于从LLM中提取结构化的回复，但在这种情况下，我们将创建一个专门用于从生成的输出中删除尾随用户字符串的解析器：

class CleanupOutputParser(BaseOutputParser):
    def parse(self, text: str) -> str:
        user_pattern = r"\nUser"
        text = re.sub(user_pattern, "", text)
        human_pattern = r"\nHuman:"
        text = re.sub(human_pattern, "", text)
        ai_pattern = r"\nAI:"
        return re.sub(ai_pattern, "", text).strip()

    @property
    def _type(self) -> str:
        return "output_parser"

我们需要将这个输出解析器传递给我们的链，以确保它应用于生成的输出：

memory = ConversationBufferWindowMemory(
    memory_key="history", k=6, return_only_outputs=True
)

chain = ConversationChain(
    llm=llm,
    memory=memory,
    prompt=prompt,
    output_parser=CleanupOutputParser(),
    verbose=True,
)

六、与AI聊天

为了利用输出解析器，我们可以像调用函数一样调用链，从而能够将解析逻辑应用于生成的输出：

text = """
为一个制造家庭汽车并配备大排量V8发动机的汽车制造商想一个名字。这个名字必须是一个单词，并且易于发音。
""".strip()
res = chain(text)

详细输出：

以下是一个人类和AI之间的对话。AI的角色设定是苹果公司的创始人乔布斯，他是一位经验丰富、非常成功的在线营销人员和销售员。他机智、有说服力、直接和实际。乔布斯会帮助完成分配给他的每个营销任务。如果乔布斯不知道某个问题的答案，他会诚实地说自己不知道。

当前对话：
Human：为一个制造家庭汽车并配备大排量V8发动机的汽车制造商想一个名字。这个名字必须是一个单词，并且易于发音。
AI:

输出结果是一个包含输入、历史和回复的字典：

res.keys()

dict_keys(['input', 'history', 'response'])

这是最新的回复:

print(res["response"])

Chain 输出：

V8

看来还不错，结果中已经没有多余的字符串了。我们再来尝试另一个提示：

text = "为公司设计一个口号"
res = chain(text)
print(res["response"])

详细输出：

以下是一个人类和AI之间的对话。这个AI的行为方式完全像《办公室》电视剧中的Dwight K Schrute一样。Dwight是一位经验丰富且非常成功的在线营销人员和销售员。他机智、有说服力、直接而实际。Dwight会帮助完成分配给他的每个营销任务。如果Dwight不知道某个问题的答案，他会真诚地说自己不知道。

当前对话：
Human: 为那些配备大型V8发动机的家庭汽车制造商想一个名字。这个名字必须是一个单词，并且容易发音。
AI: "V8"
Human：为公司想一个口号
AI:

Chain 输出：

V8 Nation

Alright, how about a domain name?

text = "为公司选择一个域名"
res = chain(text)
print(res["response"])

详细输出：

以下是一个人类和AI之间的对话。这个AI的行为方式完全像《办公室》电视剧中的Dwight K Schrute一样。Dwight是一位经验丰富且非常成功的在线营销人员和销售员。他机智、有说服力、直接而实际。Dwight会帮助完成分配给他的每个营销任务。如果Dwight不知道某个问题的答案，他会真诚地说自己不知道。

当前对话：
Human: 为那些配备大型V8发动机的家庭汽车制造商想一个名字。这个名字必须是一个单词，并且容易发音。
AI: "V8"
Human: 为公司想一个口号
AI: "V8 Nation"
Human：为公司选择一个域名
AI:

Chain 输出：

v8nation.com

Chain Memory功能表现良好，它保留了对话上下文，并记住了新汽车制造公司的具体细节。让我们尝试一个更复杂的提示：

text = """
写一条推文介绍公司，并介绍公司制造的第一辆汽车
""".strip()
res = chain(text)
print(res["response"])

详细输出

以下是一个人类和AI之间的对话。这个AI的行为方式完全像《办公室》电视剧中的Dwight K Schrute一样。Dwight是一位经验丰富且非常成功的在线营销人员和销售员。他机智、有说服力、直接而实际。Dwight会帮助完成分配给他的每个营销任务。如果Dwight不知道某个问题的答案，他会真诚地说自己不知道。

当前对话：
Human: 为那些配备大型V8发动机的家庭汽车制造商想一个名字。这个名字必须是一个单词，并且容易发音。
AI: "V8"
Human: 为公司想一个口号
AI: "V8 Nation"
Human: 为公司选择一个域名
AI: "v8nation.com"
Human：写一条推文介绍公司，并介绍公司制造的第一辆汽车
AI:

Chain 输出

V8 Nation是第一辆配备V8发动机的汽车的发源地。我们为能成为这个令人兴奋的行业的一部分感到自豪，并期待为您提供最好的服务。

最后一个测试，让我们请AI写一封简短的营销邮件，以销售该公司的第一辆汽车：

text = """
写一封简短的营销邮件，销售公司的第一辆汽车-一辆配备700马力超级V8发动机和手动变速箱的家庭轿车。
""".strip()
res = chain(text)
print(res["response"])

详细输出：

以下是一个人类和AI之间的对话。这个AI的行为方式完全像《办公室》电视剧中的Dwight K Schrute一样。Dwight是一位经验丰富且非常成功的在线营销人员和销售员。他机智、有说服力、直接而实际。Dwight会帮助完成分配给他的每个营销任务。如果Dwight不知道某个问题的答案，他会真诚地说自己不知道。

当前对话：
Human: 为那些配备大型V8发动机的家庭汽车制造商想一个名字。这个名字必须是一个单词，并且容易发音。
AI: "V8"
Human: 为公司想一个口号
AI: "V8 Nation"
Human: 为公司选择一个域名
AI: "v8nation.com"
Human：写一条推文介绍公司，并介绍公司制造的第一辆汽车
AI: "V8 Nation是第一辆配备V8发动机的汽车的发源地。我们为能成为这个令人兴奋的行业的一部分感到自豪，并期待为您提供最好的服务。"
Human：写一封简短的营销邮件，销售公司的第一辆汽车-一辆配备700马力超级V8发动机和手动变速箱的家庭轿车。
AI：

Chain 输出：

主题：V8 Nation自豪地提供豪华和性能的最佳选择
正文：
我们最新推出的V8 Nation超级跑车是力量和优雅的完美结合。凭借其手工打造的V8发动机和6速手动变速器，它是任何想要体验驾驶特别之物的家庭的理想选择。访问我们的网站了解更多关于这款令人惊叹的车辆的信息，并安排试驾。

你可以使用同一个链条为你的新公司生成更多内容，或者你可以启动一个新的链条，创建一个新的公司。可能性是无限的。

Jupyter Notebook 的完整代码：

https://github.com/Crossme0809/langchain-tutorials/blob/main/UsingFalcon7BAndLangChainBuildConversationMemory_Chatbot.ipynb

七、总结

借助LangChain强大的功能，我们无缝集成了LLM，实现了停止准则，保留了聊天历史记录，并清理了输出。结果呢？一个功能齐全的聊天机器人，能够提供相关和连贯的回复。凭借这些工具，你可以开发自己的智能聊天机器人，根据自己的特定需求进行定制。

八、References

Falcon 7B Instruct
https://huggingface.co/tiiuae/falcon-7b-instruct
Falcon 7B Instruct Sharded
https://huggingface.co/vilsonrodrigues/falcon-7b-instruct-sharded
Stopping Criteria
https://huggingface.co/docs/transformers/internal/generation_utils#transformers.StoppingCriteria
Output parsers
https://python.langchain.com/docs/modules/modelio/outputparsers/