Skip to main content

RWKV-4

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

安装和设置

  • 使用pip install rwkv安装Python包
  • 使用pip install tokenizer安装tokenizer Python包
  • 下载RWKV模型并将其放置在所需的目录中
  • 下载tokens文件

使用

RWKV

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

from langchain.llms import RWKV

# 测试模型

```python

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(generate_prompt("从前有一天,"))

模型文件

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

Rwkv-4模型 -> 推荐的VRAM

RWKV VRAM
模型 | 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支持。