johnpaulbin commited on
Commit
6db75e5
·
1 Parent(s): 8ff6ba7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -5
app.py CHANGED
@@ -5,15 +5,14 @@ from huggingface_hub import hf_hub_download
5
  from pynvml import *
6
  nvmlInit()
7
  gpu_h = nvmlDeviceGetHandleByIndex(0)
8
- ctx_limit = 1536
9
- title = "RWKV-4-World-7B-v1-20230626-ctx4096"
10
-
11
  os.environ["RWKV_JIT_ON"] = '1'
12
- os.environ["RWKV_CUDA_ON"] = '1' # if '1' then use CUDA kernel for seq mode (much faster)
13
 
14
  from rwkv.model import RWKV
15
  model_path = hf_hub_download(repo_id="BlinkDL/rwkv-4-world", filename=f"{title}.pth")
16
- model = RWKV(model=model_path, strategy='cuda fp16i8 *8 -> cuda fp16')
17
  from rwkv.utils import PIPELINE, PIPELINE_ARGS
18
  pipeline = PIPELINE(model, "rwkv_vocab_v20230424")
19
 
 
5
  from pynvml import *
6
  nvmlInit()
7
  gpu_h = nvmlDeviceGetHandleByIndex(0)
8
+ ctx_limit = 512
9
+ title = "RWKV-4-World-0.4B-v1-20230529-ctx4096"
 
10
  os.environ["RWKV_JIT_ON"] = '1'
11
+ os.environ["RWKV_CUDA_ON"] = '0' # if '1' then use CUDA kernel for seq mode (much faster)
12
 
13
  from rwkv.model import RWKV
14
  model_path = hf_hub_download(repo_id="BlinkDL/rwkv-4-world", filename=f"{title}.pth")
15
+ model = RWKV(model=model_path, strategy='cpu')
16
  from rwkv.utils import PIPELINE, PIPELINE_ARGS
17
  pipeline = PIPELINE(model, "rwkv_vocab_v20230424")
18