pvduy commited on
Commit
ca3ac1a
1 Parent(s): 658eb41

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -2
app.py CHANGED
@@ -60,8 +60,13 @@ def predict(message, history, system_prompt, temperature, max_tokens):
60
 
61
  if __name__ == "__main__":
62
  args = parse_args()
63
- tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-2-12b-chat", trust_remote_code=True)
64
- model = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-2-12b-chat", trust_remote_code=True, torch_dtype=torch.bfloat16)
 
 
 
 
 
65
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
66
  model = model.to(device)
67
  gr.ChatInterface(
 
60
 
61
  if __name__ == "__main__":
62
  args = parse_args()
63
+ tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-2-12b-chat")
64
+ tokenizer = AutoTokenizer.from_pretrained('stabilityai/stablelm-2-12b-chat')
65
+ model = AutoModelForCausalLM.from_pretrained(
66
+ 'stabilityai/stablelm-2-12b-chat',
67
+ torch_dtype=torch.bfloat16,
68
+ low_cpu_mem_usage=True
69
+ )
70
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
71
  model = model.to(device)
72
  gr.ChatInterface(