Update agents.py
Browse files
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 |
-
|
|
|
|
|
|
|
|
|
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 ---
|