Tonic commited on
Commit
b5b8829
1 Parent(s): 7563a34

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -17,9 +17,10 @@ offload_folder = './model_weights'
17
  if not os.path.exists(offload_folder):
18
  os.makedirs(offload_folder)
19
 
20
- model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto')
21
- tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False,)
22
  model = model.to(torch.bfloat16)
 
23
 
24
  class OrcaChatBot:
25
  def __init__(self, model, tokenizer, system_message="You are Orca, an AI language model created by Microsoft. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior."):
 
17
  if not os.path.exists(offload_folder):
18
  os.makedirs(offload_folder)
19
 
20
+ tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
21
+ model = AutoModelForCausalLM.from_pretrained(model_name)
22
  model = model.to(torch.bfloat16)
23
+ model = model.to(device_map='auto', offload_folder=offload_folder)
24
 
25
  class OrcaChatBot:
26
  def __init__(self, model, tokenizer, system_message="You are Orca, an AI language model created by Microsoft. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior."):