akashmishra358 commited on
Commit
b8298df
·
verified ·
1 Parent(s): 9f904c9

Update model.py

Browse files

model_url = configs["model"]["generation_model"] # Fixed "genration_model" -> "generation_model"

# 3. Add authentication to model loading
self.model = AutoModelForCausalLM.from_pretrained(
model_url,
token=self.hf_token, # Added authentication
torch_dtype=torch.float16,
low_cpu_mem_usage=True,
attn_implementation="sdpa",
device_map="auto" # Better device handling


Low CPU memory changed from false to true

Files changed (1) hide show
  1. model.py +1 -1
model.py CHANGED
@@ -33,7 +33,7 @@ class RAGModel:
33
  model_url,
34
  token=self.hf_token, # Added authentication
35
  torch_dtype=torch.float16,
36
- low_cpu_mem_usage=False,
37
  attn_implementation="sdpa",
38
  device_map="auto" # Better device handling
39
  )
 
33
  model_url,
34
  token=self.hf_token, # Added authentication
35
  torch_dtype=torch.float16,
36
+ low_cpu_mem_usage=True,
37
  attn_implementation="sdpa",
38
  device_map="auto" # Better device handling
39
  )