Coool2 commited on
Commit
07ae052
·
1 Parent(s): 979a803

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +3 -3
agent.py CHANGED
@@ -77,7 +77,6 @@ def get_max_memory_config(max_memory_per_gpu):
77
  return max_memory
78
  return None
79
 
80
- max_mem = get_max_memory_config("10GIB")
81
  # Initialize models based on API availability
82
  def initialize_models(use_api_mode=False):
83
  """Initialize LLM, Code LLM, and Embed models based on mode"""
@@ -131,8 +130,9 @@ def initialize_models(use_api_mode=False):
131
  context_window: int = 32768
132
  num_output: int = 256
133
  model_name: str = "Qwen/Qwen2.5-VL-7B-Instruct"
134
-
135
- def __init__(self, device: str = "cuda"):
 
136
  self.device = device
137
  self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
138
  self.model_name, torch_dtype=torch.bfloat16, device_map="auto"
 
77
  return max_memory
78
  return None
79
 
 
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"""
 
130
  context_window: int = 32768
131
  num_output: int = 256
132
  model_name: str = "Qwen/Qwen2.5-VL-7B-Instruct"
133
+ device: str = Field(default="cuda")
134
+
135
+ def __init__(self):
136
  self.device = device
137
  self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
138
  self.model_name, torch_dtype=torch.bfloat16, device_map="auto"