Spaces:
Configuration error
Configuration error
Update agent.py
Browse files
agent.py
CHANGED
@@ -66,10 +66,17 @@ import weave
|
|
66 |
|
67 |
weave.init("gaia-llamaindex-agents")
|
68 |
|
69 |
-
|
70 |
-
|
71 |
-
|
72 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
73 |
# Initialize models based on API availability
|
74 |
def initialize_models(use_api_mode=False):
|
75 |
"""Initialize LLM, Code LLM, and Embed models based on mode"""
|
|
|
66 |
|
67 |
weave.init("gaia-llamaindex-agents")
|
68 |
|
69 |
+
def get_max_memory_config(max_memory_per_gpu):
|
70 |
+
"""Generate max_memory config for available GPUs"""
|
71 |
+
if torch.cuda.is_available():
|
72 |
+
num_gpus = torch.cuda.device_count()
|
73 |
+
max_memory = {}
|
74 |
+
for i in range(num_gpus):
|
75 |
+
max_memory[i] = max_memory_per_gpu
|
76 |
+
return max_memory
|
77 |
+
return None
|
78 |
+
|
79 |
+
max_mem = get_max_memory_config("10GIB")
|
80 |
# Initialize models based on API availability
|
81 |
def initialize_models(use_api_mode=False):
|
82 |
"""Initialize LLM, Code LLM, and Embed models based on mode"""
|