Coool2 commited on
Commit
2281e58
·
verified ·
1 Parent(s): db65858

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +9 -3
agent.py CHANGED
@@ -101,7 +101,10 @@ embed_model = HuggingFaceEmbedding(
101
  model_name="llamaindex/vdr-2b-multi-v1",
102
  device="cpu", # "mps" for mac, "cuda" for nvidia GPUs
103
  trust_remote_code=True,
104
- )
 
 
 
105
 
106
  Settings.llm = proj_llm
107
  Settings.embed_model = embed_model
@@ -261,9 +264,12 @@ class DynamicQueryEngineManager:
261
  )
262
  self.visual_reranker = ColPaliRerank(
263
  top_n=3,
264
- device = "cpu"
 
 
 
 
265
  )
266
-
267
  def postprocess_nodes(self, nodes, query_bundle):
268
  # Your exact implementation
269
  text_nodes = []
 
101
  model_name="llamaindex/vdr-2b-multi-v1",
102
  device="cpu", # "mps" for mac, "cuda" for nvidia GPUs
103
  trust_remote_code=True,
104
+ model_kwargs={
105
+ "torch_dtype": torch.float16,
106
+ "device_map": "auto" # Optional: for better GPU memory management
107
+ })
108
 
109
  Settings.llm = proj_llm
110
  Settings.embed_model = embed_model
 
264
  )
265
  self.visual_reranker = ColPaliRerank(
266
  top_n=3,
267
+ device="cpu",
268
+ model_kwargs={
269
+ "torch_dtype": torch.float16,
270
+ "device_map": "auto"
271
+ }
272
  )
 
273
  def postprocess_nodes(self, nodes, query_bundle):
274
  # Your exact implementation
275
  text_nodes = []