跳至内容

RWKV-4

本页面介绍如何在 LangChain 中使用 RWKV-4 包装器。它分为两个部分:安装和设置,以及带有示例的使用。

安装和设置

  • 使用 pip install rwkv 安装 Python 包

  • 使用 pip install tokenizer 安装分词器 Python 包

  • 下载一个 RWKV 模型 并将其放置在所需的目录中

  • 下载 tokens 文件

用法

RWKV

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

from langchain.llms import RWKV

# Test the model

def generate_prompt(instruction, input=None):
    if input:
        return f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

# Instruction:
{instruction}

# Input:
{input}

# Response:
"""
    else:
        return f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.

# Instruction:
{instruction}

# Response:
"""

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

模型文件

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

Rwkv-4 models -> 推荐 VRAM

RWKV VRAM
Model | 8bit | bf16/fp16 | fp32
14B   | 16GB | 28GB      | >50GB
7B    | 8GB  | 14GB      | 28GB
3B    | 2.8GB| 6GB       | 12GB
1b5   | 1.3GB| 3GB       | 6GB

查看 rwkv pip 页面获取更多关于策略的信息,包括流处理和 cuda 支持。