budivoy commited on
Commit
64ed80a
·
verified ·
1 Parent(s): 3870827

Update agents.py

Browse files
Files changed (1) hide show
  1. agents.py +3 -3
agents.py CHANGED
@@ -168,11 +168,11 @@ def build_graph(provider: str = "rwkv"):
168
  )
169
  elif provider == "rwkv":
170
  # --- BEGIN RWKV SETUP ---
171
- title = "rwkv7-g1-2.9b-20250519-ctx4096"
172
  pth = hf_hub_download(repo_id="BlinkDL/rwkv7-g1", filename=f"{title}.pth")
173
  model_path = pth.replace(".pth", "")
174
 
175
- raw_llm = RWKV(model=model_path, strategy='cpu fp32')
176
  pipeline = PIPELINE(raw_llm, "rwkv_vocab_v20230424")
177
 
178
  class RWKVWithTools:
@@ -211,7 +211,7 @@ def build_graph(provider: str = "rwkv"):
211
  print(f'Prompt: {prompt}')
212
 
213
  # delegate to RWKV invoke()
214
- out_str = self.pipeline.generate(prompt, token_count=1000)
215
 
216
  print(f'Response: {out_str}')
217
 
 
168
  )
169
  elif provider == "rwkv":
170
  # --- BEGIN RWKV SETUP ---
171
+ title = "rwkv7-g1-1.5b-20250429-ctx4096"
172
  pth = hf_hub_download(repo_id="BlinkDL/rwkv7-g1", filename=f"{title}.pth")
173
  model_path = pth.replace(".pth", "")
174
 
175
+ raw_llm = RWKV(model=model_path, strategy='cpu fp16')
176
  pipeline = PIPELINE(raw_llm, "rwkv_vocab_v20230424")
177
 
178
  class RWKVWithTools:
 
211
  print(f'Prompt: {prompt}')
212
 
213
  # delegate to RWKV invoke()
214
+ out_str = self.pipeline.generate(prompt, token_count=300)
215
 
216
  print(f'Response: {out_str}')
217