RootYuan commited on
Commit
4a48591
·
1 Parent(s): 120c0ee

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -3
app.py CHANGED
@@ -1,4 +1,3 @@
1
- import torch
2
  import gradio as gr
3
  from uuid import uuid4
4
  from threading import Thread
@@ -9,7 +8,7 @@ from theme_dropdown import create_theme_dropdown
9
 
10
  model_name = "RootYuan/RootYuan-RedLing-7B-v0.1"
11
  max_new_tokens = 2048
12
- device = 'cuda' if torch.cuda.is_available() else 'cpu'
13
 
14
 
15
  DEFAULT_SYSTEM_MESSAGE = """
@@ -23,7 +22,7 @@ EOT_TOKEN = "<EOT>"
23
  PROMPT_TEMPLATE = "USER:{user}<EOT>ASSISTANT:{assistant}{eos_token}"
24
 
25
  tokenizer = AutoTokenizer.from_pretrained(model_name)
26
- model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
27
 
28
  dropdown, js = create_theme_dropdown()
29
 
 
 
1
  import gradio as gr
2
  from uuid import uuid4
3
  from threading import Thread
 
8
 
9
  model_name = "RootYuan/RootYuan-RedLing-7B-v0.1"
10
  max_new_tokens = 2048
11
+ device = 'cpu'
12
 
13
 
14
  DEFAULT_SYSTEM_MESSAGE = """
 
22
  PROMPT_TEMPLATE = "USER:{user}<EOT>ASSISTANT:{assistant}{eos_token}"
23
 
24
  tokenizer = AutoTokenizer.from_pretrained(model_name)
25
+ model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
26
 
27
  dropdown, js = create_theme_dropdown()
28