BenBranyon commited on
Commit
50f12a0
·
verified ·
1 Parent(s): f761b13

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -4
app.py CHANGED
@@ -7,10 +7,11 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
7
  client = InferenceClient("BenBranyon/zephyr-sumbot-all-songs")
8
 
9
  #Transformers Code
10
- model_id = "BenBranyon/zephyr-sumbot-all-songs"
11
- model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_4bit=True)
12
- tokenizer = AutoTokenizer.from_pretrained(model_id)
13
- tokenizer.use_default_system_prompt = False
 
14
 
15
 
16
  def respond(
 
7
  client = InferenceClient("BenBranyon/zephyr-sumbot-all-songs")
8
 
9
  #Transformers Code
10
+ if torch.cuda.is_available():
11
+ model_id = "BenBranyon/zephyr-sumbot-all-songs"
12
+ model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_4bit=True)
13
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
14
+ tokenizer.use_default_system_prompt = False
15
 
16
 
17
  def respond(