Skip to main content
Open on GitHub

RWKV-4

此页面介绍如何在 LangChain 中使用 RWKV-4 包装器。 分为两部分：安装与设置，然后是实际使用示例。

安装与设置​

  • 使用 pip install rwkv 安装 Python 包
  • 使用 pip install tokenizer 安装分词器 Python 包
  • 下载一个 RWKV 模型 并将其放置在您想要的目录中
  • 下载 tokens 文件

用法​

RWKV​

要使用 RWKV 包装器，您需要提供预训练模型文件的路径和分词器的配置。

from langchain_community.llms import RWKV

# 测试模型

API Reference:RWKV
def generate_prompt(instruction, input=None):
    if input:
        return f"""下面是一个描述任务的指令，以及提供进一步上下文的输入。请写一个适当完成请求的响应。

# 指令:
{instruction}

# 输入:
{input}

# 响应:
"""
    else:
        return f"""下面是一个描述任务的指令。请写一个适当完成请求的响应。

# 指令:
{instruction}

# 响应:
"""


model = RWKV(model="./models/RWKV-4-Raven-3B-v7-Eng-20230404-ctx4096.pth", strategy="cpu fp32", tokens_path="./rwkv/20B_tokenizer.json")
response = model.invoke(generate_prompt("Once upon a time, "))

模型文件​

您可以在 RWKV-4-Raven 仓库中找到模型文件下载链接。

Rwkv-4 模型 -> 推荐显存​

RWKV 显存
模型 | 8bit | bf16/fp16 | fp32
14B   | 16GB | 28GB      | >50GB
7B    | 8GB  | 14GB      | 28GB
3B    | 2.8GB| 6GB       | 12GB
1b5   | 1.3GB| 3GB       | 6GB

有关策略（包括流式传输和 CUDA 支持）的更多信息，请参阅 rwkv pip 页面。