budivoy commited on
Commit
dc20980
·
verified ·
1 Parent(s): 60e2f92

Update agents.py

Browse files
Files changed (1) hide show
  1. agents.py +11 -1
agents.py CHANGED
@@ -14,6 +14,10 @@ from langchain_community.vectorstores import SupabaseVectorStore
14
  from langchain_core.messages import SystemMessage, HumanMessage
15
  from langchain_core.tools import tool
16
 
 
 
 
 
17
  from huggingface_hub import hf_hub_download
18
  from rwkv.model import RWKV
19
  from rwkv.utils import PIPELINE, PIPELINE_ARGS
@@ -203,9 +207,15 @@ def build_graph(provider: str = "rwkv"):
203
  )
204
 
205
  prompt = header + convo
 
 
206
 
207
  # delegate to RWKV invoke()
208
- return self.pipeline.generate(prompt, token_count=200)
 
 
 
 
209
 
210
  llm = RWKVWithTools(raw_llm, system_prompt=system_prompt)
211
  # --- END RWKV SETUP ---
 
14
  from langchain_core.messages import SystemMessage, HumanMessage
15
  from langchain_core.tools import tool
16
 
17
+ os.environ["RWKV_V7_ON"] = '1' # ==> enable RWKV-7 mode
18
+ os.environ['RWKV_JIT_ON'] = '1' # '1' for better speed
19
+ os.environ["RWKV_CUDA_ON"] = '0' # '1' to compile CUDA kernel (10x faster), requires c++ compiler & cuda libraries
20
+
21
  from huggingface_hub import hf_hub_download
22
  from rwkv.model import RWKV
23
  from rwkv.utils import PIPELINE, PIPELINE_ARGS
 
207
  )
208
 
209
  prompt = header + convo
210
+
211
+ print(f'Prompt: {prompt}')
212
 
213
  # delegate to RWKV invoke()
214
+ out_str = self.pipeline.generate(prompt, token_count=200)
215
+
216
+ print(f'Response: {out_str}')
217
+
218
+ return out_str
219
 
220
  llm = RWKVWithTools(raw_llm, system_prompt=system_prompt)
221
  # --- END RWKV SETUP ---