ejschwartz commited on
Commit
b1dd808
·
1 Parent(s): 516424d

format stuff

Browse files
Files changed (1) hide show
  1. app.py +12 -9
app.py CHANGED
@@ -2,26 +2,29 @@ import gradio as gr
2
  import spaces
3
  import torch
4
  import os
5
-
6
  from transformers import AutoTokenizer, AutoModelForCausalLM
7
- from huggingface_hub import login
8
 
9
- hf_key = os.environ['HF_TOKEN']
10
- login(token = hf_key)
 
 
11
 
12
- tokenizer = AutoTokenizer.from_pretrained('bigcode/starcoderbase-3b', use_auth_token=hf_key)
 
 
13
  vardecoder_model = AutoModelForCausalLM.from_pretrained(
14
- "ejschwartz/resym-vardecoder",
15
- torch_dtype=torch.bfloat16, device_map='auto'
16
  )
17
 
18
  zero = torch.Tensor([0]).cuda()
19
- print(zero.device) # <-- 'cpu' 🤔
 
20
 
21
  @spaces.GPU
22
  def greet(n):
23
- print(zero.device) # <-- 'cuda:0' 🤗
24
  return f"Hello {zero + n} Tensor"
25
 
 
26
  demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
27
  demo.launch()
 
2
  import spaces
3
  import torch
4
  import os
 
5
  from transformers import AutoTokenizer, AutoModelForCausalLM
 
6
 
7
+ import huggingface_hub
8
+
9
+ hf_key = os.environ["HF_TOKEN"]
10
+ huggingface_hub.login(token=hf_key)
11
 
12
+ tokenizer = AutoTokenizer.from_pretrained(
13
+ "bigcode/starcoderbase-3b", use_auth_token=hf_key
14
+ )
15
  vardecoder_model = AutoModelForCausalLM.from_pretrained(
16
+ "ejschwartz/resym-vardecoder", torch_dtype=torch.bfloat16, device_map="auto"
 
17
  )
18
 
19
  zero = torch.Tensor([0]).cuda()
20
+ print(zero.device) # <-- 'cpu' 🤔
21
+
22
 
23
  @spaces.GPU
24
  def greet(n):
25
+ print(zero.device) # <-- 'cuda:0' 🤗
26
  return f"Hello {zero + n} Tensor"
27
 
28
+
29
  demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
30
  demo.launch()