Spaces:
Configuration error
Configuration error
Update agent.py
Browse files
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 |
-
|
|
|
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"
|