Update agent.py
Browse files
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
|
|
|
|
|
|
|
|
|
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 = []
|