Run Huggingface RWKV5 World Model

This model is developed and converted through https://github.com/BBuf/RWKV-World-HF-Tokenizer. If you have any issues, you can raise them in this project. You're also welcome to star it to follow the subsequent development progress.

CPU

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("RWKV/rwkv-5-world-169m", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-5-world-169m", trust_remote_code=True)

text = "\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese."
prompt = f'Question: {text.strip()}\n\nAnswer:'

inputs = tokenizer(prompt, return_tensors="pt")
output = model.generate(inputs["input_ids"], max_new_tokens=256)
print(tokenizer.decode(output[0].tolist(), skip_special_tokens=True))

output:

Question: In a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese.

Answer: The scientists found a herd of dragons living in a remote, previously unexplored valley, in Tibet.

GPU

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("RWKV/rwkv-5-world-169m", trust_remote_code=True, torch_dtype=torch.float16).to(0)
tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-5-world-169m", trust_remote_code=True)

text = "你叫什么名字？"
prompt = f'Question: {text.strip()}\n\nAnswer:'

inputs = tokenizer(prompt, return_tensors="pt").to(0)
output = model.generate(inputs["input_ids"], max_new_tokens=256, do_sample=True, temperature=1.0, top_p=0.1, top_k=0)
print(tokenizer.decode(output[0].tolist(), skip_special_tokens=True))

output:

Question: 你叫什么名字？

Answer: 我叫李明，是一名程序员。