test1.py

import os
os.environ['RWKV_JIT_ON'] = '1'
os.environ["RWKV_CUDA_ON"] = '0'

from rwkv.model import RWKV
from rwkv.utils import PIPELINE, PIPELINE_ARGS

model = RWKV(model='D:\\my_files\\RWKV-x060-World-1B6-v2.1-20240328-ctx4096', strategy='cpu fp32')
pipeline = PIPELINE(model, "rwkv_vocab_v20230424")

args = PIPELINE_ARGS(temperature = 0.5, top_p = 0.7, top_k = 100, # top_k = 0 then ignore
                     alpha_frequency = 0.25,
                     alpha_presence = 0.25,
                     alpha_decay = 0.996, # gradually decay the penalty
                     chunk_len = 256) # split input into chunks to save VRAM (shorter -> slower)

inputString = '''Instruction: 按以下内容用英文描述画面

Input: 一只熊在喜马拉雅山上吃鱼

Response:'''

for _ in range(10):
    print(pipeline.generate(inputString, token_count=200, args=args))