ejschwartz commited on
Commit
d3eb07d
·
1 Parent(s): 01c3073

try to download models

Browse files
Files changed (1) hide show
  1. app.py +13 -0
app.py CHANGED
@@ -2,6 +2,19 @@ import gradio as gr
2
  import spaces
3
  import torch
4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5
  zero = torch.Tensor([0]).cuda()
6
  print(zero.device) # <-- 'cpu' 🤔
7
 
 
2
  import spaces
3
  import torch
4
 
5
+ import torch
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
+ model = AutoModelForCausalLM.from_pretrained(
14
+ model_path, use_auth_token=hf_key,
15
+ torch_dtype=torch.bfloat16, device_map='auto'
16
+ )
17
+
18
  zero = torch.Tensor([0]).cuda()
19
  print(zero.device) # <-- 'cpu' 🤔
20