Coool2 commited on
Commit
5aecaaa
·
1 Parent(s): b1f5ba8

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +11 -4
agent.py CHANGED
@@ -66,10 +66,17 @@ import weave
66
 
67
  weave.init("gaia-llamaindex-agents")
68
 
69
- from transformers import get_max_memory
70
-
71
- max_mem = get_max_memory(0.75) # 75% of each device's memory
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"""