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

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +7 -6
agent.py CHANGED
@@ -99,12 +99,13 @@ code_llm = HuggingFaceLLM(
99
 
100
  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
  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
@@ -266,8 +267,8 @@ class DynamicQueryEngineManager:
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):
 
99
 
100
  embed_model = HuggingFaceEmbedding(
101
  model_name="llamaindex/vdr-2b-multi-v1",
102
+ device="cpu",
103
  trust_remote_code=True,
104
  model_kwargs={
105
+ "torch_dtype": torch.float32, # Use float32 for CPU
106
+ "low_cpu_mem_usage": True, # Still get memory optimization
107
+ }
108
+ )
109
 
110
  Settings.llm = proj_llm
111
  Settings.embed_model = embed_model
 
267
  top_n=3,
268
  device="cpu",
269
  model_kwargs={
270
+ "torch_dtype": torch.float32, # Use float32 for CPU
271
+ "low_cpu_mem_usage": True, # Still get memory optimization
272
  }
273
  )
274
  def postprocess_nodes(self, nodes, query_bundle):