|
@@ -13,7 +13,6 @@ np.set_printoptions(precision=4, suppress=True, linewidth=200)
|
|
|
os.environ['RWKV_JIT_ON'] = '1'
|
|
os.environ['RWKV_JIT_ON'] = '1'
|
|
|
os.environ["RWKV_CUDA_ON"] = '0' # '1' : use CUDA kernel for seq mode (much faster)
|
|
os.environ["RWKV_CUDA_ON"] = '0' # '1' : use CUDA kernel for seq mode (much faster)
|
|
|
|
|
|
|
|
-import repositories.ChatRWKV.v2.rwkv as rwkv
|
|
|
|
|
from rwkv.model import RWKV
|
|
from rwkv.model import RWKV
|
|
|
from rwkv.utils import PIPELINE, PIPELINE_ARGS
|
|
from rwkv.utils import PIPELINE, PIPELINE_ARGS
|
|
|
|
|
|
|
@@ -22,6 +21,6 @@ def load_RWKV_model(path):
|
|
|
print(f'strategy={"cpu" if shared.args.cpu else "cuda"} {"fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16"}')
|
|
print(f'strategy={"cpu" if shared.args.cpu else "cuda"} {"fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16"}')
|
|
|
|
|
|
|
|
model = RWKV(model=path.as_posix(), strategy=f'{"cpu" if shared.args.cpu else "cuda"} {"fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16"}')
|
|
model = RWKV(model=path.as_posix(), strategy=f'{"cpu" if shared.args.cpu else "cuda"} {"fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16"}')
|
|
|
- pipeline = PIPELINE(model, Path("repositories/ChatRWKV/20B_tokenizer.json").as_posix())
|
|
|
|
|
|
|
+ pipeline = PIPELINE(model, Path("models/20B_tokenizer.json").as_posix())
|
|
|
|
|
|
|
|
return pipeline
|
|
return pipeline
|