DaniilAlpha commited on
Commit
15862bc
·
1 Parent(s): 74515cd

Update answerer.py

Browse files
Files changed (1) hide show
  1. answerer.py +1 -10
answerer.py CHANGED
@@ -4,21 +4,12 @@ from huggingface_hub import hf_hub_download
4
  from rwkv.model import RWKV
5
  from rwkv.utils import PIPELINE, PIPELINE_ARGS
6
 
7
- ### settings ###
8
-
9
- ###
10
-
11
-
12
-
13
- os.environ["RWKV_JIT_ON"] = "1"
14
- # os.environ["RWKV_CUDA_ON"] = "1" # if "1" then use CUDA kernel for seq mode (much faster)
15
-
16
  class Answerer:
17
  def __init__(self, repo: str, filename: str, vocab: str, strategy: str, ctx_limit: int):
18
  os.environ["RWKV_JIT_ON"] = "1"
19
  # os.environ["RWKV_CUDA_ON"] = "1"
20
 
21
- self.__model = RWKV(hf_hub_download(repo, f"{filename}.pth"), strategy=strategy)
22
  self.__pipeline = PIPELINE(self.__model, vocab)
23
  self.ctx_limit = ctx_limit
24
 
 
4
  from rwkv.model import RWKV
5
  from rwkv.utils import PIPELINE, PIPELINE_ARGS
6
 
 
 
 
 
 
 
 
 
 
7
  class Answerer:
8
  def __init__(self, repo: str, filename: str, vocab: str, strategy: str, ctx_limit: int):
9
  os.environ["RWKV_JIT_ON"] = "1"
10
  # os.environ["RWKV_CUDA_ON"] = "1"
11
 
12
+ self.__model = RWKV(hf_hub_download(repo, filename), strategy=strategy)
13
  self.__pipeline = PIPELINE(self.__model, vocab)
14
  self.ctx_limit = ctx_limit
15