Coool2 commited on
Commit
a1ca93d
·
verified ·
1 Parent(s): 61f82e6

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +3 -3
agent.py CHANGED
@@ -78,7 +78,7 @@ proj_llm = HuggingFaceLLM(
78
  model_name=model_id,
79
  tokenizer_name=model_id,
80
  device_map="auto",
81
- model_kwargs={"torch_dtype": torch.float16},
82
  generate_kwargs={"temperature": 0.1, "top_p": 0.3} # More focused
83
  )
84
 
@@ -87,7 +87,7 @@ code_llm = HuggingFaceLLM(
87
  tokenizer_name="Qwen/Qwen2.5-Coder-3B-Instruct",
88
  device_map= "cpu", # Specify device here instead
89
  model_kwargs={
90
- "torch_dtype": torch.float32, # Use float32 for CPU
91
  "low_cpu_mem_usage": True, # Memory optimization
92
  },
93
  # Set generation parameters for precise, non-creative code output
@@ -99,7 +99,7 @@ embed_model = HuggingFaceEmbedding(
99
  device="cpu",
100
  trust_remote_code=True,
101
  model_kwargs={
102
- "torch_dtype": torch.float32, # Use float32 for CPU
103
  "low_cpu_mem_usage": True, # Still get memory optimization
104
  }
105
  )
 
78
  model_name=model_id,
79
  tokenizer_name=model_id,
80
  device_map="auto",
81
+ model_kwargs={"torch_dtype": "auto"},
82
  generate_kwargs={"temperature": 0.1, "top_p": 0.3} # More focused
83
  )
84
 
 
87
  tokenizer_name="Qwen/Qwen2.5-Coder-3B-Instruct",
88
  device_map= "cpu", # Specify device here instead
89
  model_kwargs={
90
+ "torch_dtype": "auto", # Use float32 for CPU
91
  "low_cpu_mem_usage": True, # Memory optimization
92
  },
93
  # Set generation parameters for precise, non-creative code output
 
99
  device="cpu",
100
  trust_remote_code=True,
101
  model_kwargs={
102
+ "torch_dtype": "auto", # Use float32 for CPU
103
  "low_cpu_mem_usage": True, # Still get memory optimization
104
  }
105
  )