aixsatoshi commited on
Commit
35f258a
·
verified ·
1 Parent(s): a490149

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -6
app.py CHANGED
@@ -6,14 +6,13 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
6
  import gradio as gr
7
  from threading import Thread
8
 
 
9
  model_id = "ibm-granite/granite-3.0-8b-instruct"
10
  tokenizer = AutoTokenizer.from_pretrained(model_id)
11
- model = AutoModelForCausalLM.from_pretrained(
12
- model_id,
13
- torch_dtype=torch.bfloat16,
14
- device_map="auto",
15
- #use_flash_attention_2=True,
16
- )
17
 
18
  TITLE = "<h1><center>ibm-granite/granite-3.0-8b-instruct Chat webui</center></h1>"
19
 
 
6
  import gradio as gr
7
  from threading import Thread
8
 
9
+ device = "auto"
10
  model_id = "ibm-granite/granite-3.0-8b-instruct"
11
  tokenizer = AutoTokenizer.from_pretrained(model_id)
12
+ # drop device_map if running on CPU
13
+ model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
14
+ model.eval()
15
+ # change input text as desired
 
 
16
 
17
  TITLE = "<h1><center>ibm-granite/granite-3.0-8b-instruct Chat webui</center></h1>"
18