FastChat:打造你自己的聊天机器人

文摘   2024-11-09 08:02   上海  

大家好!我是风哥,一个专注AI领域的Python工程师。今天给大家介绍一个超级强大的开源聊天机器人框架 —— FastChat!它能让你轻松部署自己的AI助手,快来一起学习吧!

基础安装与配置

首先,安装FastChat:

1pip install fschat

如果想要完整功能,建议这样安装:

1pip install "fschat[model_worker,webui]"

快速启动服务

FastChat提供了三个主要组件,需要分别启动:

1# 1. 控制器
2python -m fastchat.serve.controller
3
4# 2. 模型Worker
5python -m fastchat.serve.model_worker --model-path lmsys/vicuna-7b-v1.5
6
7# 3. Web界面
8python -m fastchat.serve.gradio_web_server

核心功能实现

1. 基础对话实现

 1from fastchat.model import load_model, get_conversation_template
2
3# 加载模型
4model, tokenizer = load_model("lmsys/vicuna-7b-v1.5")
5
6# 创建对话模板
7conv = get_conversation_template("vicuna")
8conv.append_message(conv.roles[0], "你好!")
9conv.append_message(conv.roles[1], None)
10
11# 生成回复
12prompt = conv.get_prompt()
13output = model.generate(prompt, max_length=2048)
14response = tokenizer.decode(output)
15print(response)

2. 自定义对话设置

 1from fastchat.serve.inference import ChatIO, chat_loop
2
3class CustomChatIO(ChatIO):
4    def prompt_for_input(self, role) -> str:
5        return input(f"{role}: ")
6
7    def prompt_for_output(self, role: str):
8        print(f"{role}: ", end="", flush=True)
9
10    def stream_output(self, output_text: str):
11        print(output_text, end="", flush=True)
12
13# 启动对话循环
14chat_loop(
15    model_path="lmsys/vicuna-7b-v1.5",
16    temperature=0.7,
17    max_new_tokens=512,
18    chatio=CustomChatIO(),
19)

高级特性

1. 多轮对话管理

 1from fastchat.conversation import Conversation
2
3def create_chat_session():
4    conv = Conversation(
5        system="你是一个友好的AI助手。",
6        roles=("用户""助手"),
7        messages=[],
8        offset=0,
9        sep_style=SeparatorStyle.TWO,
10    )
11
12    # 添加对话历史
13    conv.append_message("用户""今天天气真好!")
14    conv.append_message("助手""是的,春暖花开的季节最适合出门走走。")
15
16    return conv

2. 参数调优

 1def generate_response(prompt, model, tokenizer, **kwargs):
2    # 自定义生成参数
3    generation_config = {
4        "temperature"0.7,  # 温度值,控制随机性
5        "top_p"0.9,       # 核采样阈值
6        "top_k"40,        # 保留概率最高的k个token
7        "num_beams"4,     # beam search的束宽
8        "max_new_tokens"512,  # 最大生成长度
9        "repetition_penalty"1.1,  # 重复惩罚项
10    }
11
12    # 更新自定义参数
13    generation_config.update(kwargs)
14
15    # 生成回复
16    output_ids = model.generate(
17        tokenizer(prompt, return_tensors="pt").input_ids,
18        **generation_config
19    )
20
21    return tokenizer.decode(output_ids[0])

3. 流式输出实现

 1import time
2
3class StreamingChatIO(ChatIO):
4    def stream_output(self, output_text: str):
5        for char in output_text:
6            print(char, end="", flush=True)
7            time.sleep(0.02)  # 模拟打字效果
8        print()
9
10# 使用流式输出
11chat_loop(
12    model_path="lmsys/vicuna-7b-v1.5",
13    chatio=StreamingChatIO(),
14)

4. 自定义模型加载

 1def load_custom_model():
2    # 自定义模型配置
3    model_config = {
4        "model_path""your_model_path",
5        "device""cuda",  # 或 "cpu"
6        "num_gpus"1,
7        "max_gpu_memory""15GiB",
8        "load_8bit"False,  # 8位量化
9        "cpu_offloading"False,
10    }
11
12    # 加载模型
13    model = load_model(
14        model_config["model_path"],
15        device=model_config["device"],
16        num_gpus=model_config["num_gpus"],
17        max_gpu_memory=model_config["max_gpu_memory"],
18        load_8bit=model_config["load_8bit"],
19    )
20
21    return model

实用技巧

1. 异常处理

1def safe_chat_response(prompt, model, tokenizer):
2    try:
3        response = generate_response(prompt, model, tokenizer)
4        return response
5    except Exception as e:
6        print(f"生成回复时出错: {str(e)}")
7        return "抱歉,我现在无法正常回答,请稍后再试。"

2. 对话历史保存

 1import json
2
3def save_conversation(conv, filename):
4    history = {
5        "messages": conv.messages,
6        "system": conv.system,
7        "roles": conv.roles,
8    }
9
10    with open(filename, 'w', encoding='utf-8'as f:
11        json.dump(history, f, ensure_ascii=False, indent=2)
12
13def load_conversation(filename):
14    with open(filename, 'r', encoding='utf-8'as f:
15        history = json.load(f)
16
17    conv = Conversation(
18        system=history["system"],
19        roles=history["roles"],
20        messages=history["messages"],
21    )
22    return conv

今天的Python AI开发分享就到这里啦!FastChat是个功能强大的框架,还有很多有趣的特性等待大家探索。有问题随时在评论区问我,记得点赞收藏,下次给大家带来更多AI开发技巧!

py学习基地ai
分享生活百态,情感故事,了解不一样的人生
 最新文章