Mohammed-Altaf commited on
Commit
da50ade
1 Parent(s): cf4ea8b

added accelerate to the requirements

Browse files
Files changed (2) hide show
  1. load_meta_data.py +1 -0
  2. requirements.txt +2 -0
load_meta_data.py CHANGED
@@ -16,6 +16,7 @@ class ChatBot:
16
  self.tokenizer = AutoTokenizer.from_pretrained(model_id,
17
  device_map = 'auto')
18
  self.model = AutoModelForCausalLM.from_pretrained(model_id,
 
19
  device_map='auto')
20
 
21
  def get_response(self,text: UserQuery) -> str:
 
16
  self.tokenizer = AutoTokenizer.from_pretrained(model_id,
17
  device_map = 'auto')
18
  self.model = AutoModelForCausalLM.from_pretrained(model_id,
19
+ low_cpu_mem_usage=True,
20
  device_map='auto')
21
 
22
  def get_response(self,text: UserQuery) -> str:
requirements.txt CHANGED
@@ -1,4 +1,6 @@
1
  transformers
 
 
2
  fastapi
3
  torch
4
  uvicorn
 
1
  transformers
2
+ accelerate
3
+ bitsandbytes
4
  fastapi
5
  torch
6
  uvicorn