Coool2 commited on
Commit
6b13586
·
1 Parent(s): 0b15a3d

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +2 -2
agent.py CHANGED
@@ -200,14 +200,14 @@ def initialize_models(use_api_mode=False):
200
 
201
  embed_model = HuggingFaceEmbedding(
202
  model_name="llamaindex/vdr-2b-multi-v1",
203
- device="cuda:1",
204
  trust_remote_code = True)
205
 
206
  # Code LLM
207
  code_llm = HuggingFaceLLM(
208
  model_name="Qwen/Qwen2.5-Coder-1.5B-Instruct",
209
  tokenizer_name="Qwen/Qwen2.5-Coder-1.5B-Instruct",
210
- device_map="cuda:0",
211
  model_kwargs={"torch_dtype": "auto"},
212
  generate_kwargs={"do_sample": False}
213
  )
 
200
 
201
  embed_model = HuggingFaceEmbedding(
202
  model_name="llamaindex/vdr-2b-multi-v1",
203
+ device="auto",
204
  trust_remote_code = True)
205
 
206
  # Code LLM
207
  code_llm = HuggingFaceLLM(
208
  model_name="Qwen/Qwen2.5-Coder-1.5B-Instruct",
209
  tokenizer_name="Qwen/Qwen2.5-Coder-1.5B-Instruct",
210
+ device_map="auto",
211
  model_kwargs={"torch_dtype": "auto"},
212
  generate_kwargs={"do_sample": False}
213
  )