Coool2 commited on
Commit
65390be
·
verified ·
1 Parent(s): 3b576d3

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +8 -1
agent.py CHANGED
@@ -18,13 +18,20 @@ from llama_index.callbacks.wandb import WandbCallbackHandler
18
  from llama_index.core.callbacks.base import CallbackManager
19
  from llama_index.core.callbacks.llama_debug import LlamaDebugHandler
20
  from llama_index.core import Settings
 
21
 
22
  from transformers import AutoModelForCausalLM, AutoTokenizer
23
  from llama_index.llms.huggingface import HuggingFaceLLM
24
  import requests
25
 
 
26
  llama_debug = LlamaDebugHandler(print_trace_on_end=True)
27
- callback_manager = CallbackManager([llama_debug])
 
 
 
 
 
28
 
29
 
30
  model_id = "Qwen/Qwen2.5-7B-Instruct"
 
18
  from llama_index.core.callbacks.base import CallbackManager
19
  from llama_index.core.callbacks.llama_debug import LlamaDebugHandler
20
  from llama_index.core import Settings
21
+ from llama_index.callbacks.wandb import WandbCallbackHandler
22
 
23
  from transformers import AutoModelForCausalLM, AutoTokenizer
24
  from llama_index.llms.huggingface import HuggingFaceLLM
25
  import requests
26
 
27
+ wandb_callback = WandbCallbackHandler(run_args={"project": "gaia-llamaindex-agents"})
28
  llama_debug = LlamaDebugHandler(print_trace_on_end=True)
29
+
30
+ # Comprehensive callback manager
31
+ callback_manager = CallbackManager([
32
+ wandb_callback, # For W&B tracking
33
+ llama_debug # For general debugging
34
+ ])
35
 
36
 
37
  model_id = "Qwen/Qwen2.5-7B-Instruct"