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