rwkv

import os import torch from rwkv.model import RWKV from rwkv.utils import PIPELINE, PIPELINE_ARGS # Use environment variable for model path, with fallback model_path = os.environ.get('RWKV_MODEL_PATH', './RWKV-4-World-0.4B-v1-20230806-ctx4096.pth') tokenizer_path = os.environ.get('RWKV_TOKENIZER_PATH', './20B_tokenizer.json') model = RWKV(model=model_path, strategy='cuda fp16') pipeline = PIPELINE(model, tokenizer_path) # Run inference context = "\nThe universe is" state = None out_tokens = [] for i in range(100): tokens = pipeline.encode(context) if i == 0 else [token] out, state = model.forward(tokens, state) token = pipeline.sample_logits(out, temperature=1.0, top_p=0.85) out_tokens += [token] context = pipeline.decode(out_tokens) print(context)