dar-tau commited on
Commit
25fb83c
1 Parent(s): 9b0b00b

Update configs.py

Browse files
Files changed (1) hide show
  1. configs.py +3 -3
configs.py CHANGED
@@ -19,9 +19,9 @@ model_info = {
19
  interpretation_prompt_template='<s>[INST] [X] [/INST] {prompt}',
20
  layers_format=llama_layers_format),
21
  'LLAMA2-13B': dict(model_path='meta-llama/Llama-2-13b-chat-hf',
22
- token=os.environ['hf_token'], device_map='auto',
23
- max_memory={0: "15GB", 1: "30GB"}, # load_in_8bit=True,
24
- original_prompt_template='<s>{prompt}', dont_cuda=True,
25
  interpretation_prompt_template='<s>[INST] [X] [/INST] {prompt}',
26
  layers_format=llama_layers_format),
27
  'GPT-2 Small': dict(model_path='gpt2', original_prompt_template='{prompt}',
 
19
  interpretation_prompt_template='<s>[INST] [X] [/INST] {prompt}',
20
  layers_format=llama_layers_format),
21
  'LLAMA2-13B': dict(model_path='meta-llama/Llama-2-13b-chat-hf',
22
+ token=os.environ['hf_token'], torch_dtype=torch.float16,
23
+ # device_map='auto', max_memory={0: "15GB", 1: "30GB"}, dont_cuda=True, # load_in_8bit=True,
24
+ original_prompt_template='<s>{prompt}',
25
  interpretation_prompt_template='<s>[INST] [X] [/INST] {prompt}',
26
  layers_format=llama_layers_format),
27
  'GPT-2 Small': dict(model_path='gpt2', original_prompt_template='{prompt}',