liujch1998 commited on
Commit
c5c065e
Β·
1 Parent(s): 5d92357

Add auth token

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -1,4 +1,5 @@
1
  import gradio as gr
 
2
  import torch
3
  import transformers
4
 
@@ -6,8 +7,8 @@ device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cp
6
 
7
  class Interactive:
8
  def __init__(self):
9
- self.tokenizer = transformers.AutoTokenizer.from_pretrained('liujch1998/cd-pi')
10
- self.model = transformers.AutoModelForSeq2SeqLM.from_pretrained('liujch1998/cd-pi').to(device)
11
  self.linear = torch.nn.Linear(self.model.shared.embedding_dim, 1).to(device)
12
  self.linear.weight = torch.nn.Parameter(self.model.shared.weight[32099, :].unsqueeze(0)) # (1, D)
13
  self.linear.bias = torch.nn.Parameter(self.model.shared.weight[32098, 0].unsqueeze(0)) # (1)
 
1
  import gradio as gr
2
+ import os
3
  import torch
4
  import transformers
5
 
 
7
 
8
  class Interactive:
9
  def __init__(self):
10
+ self.tokenizer = transformers.AutoTokenizer.from_pretrained('liujch1998/cd-pi', use_auth_token=os.environ['HF_TOKEN_DOWNLOAD'])
11
+ self.model = transformers.AutoModelForSeq2SeqLM.from_pretrained('liujch1998/cd-pi', use_auth_token=os.environ['HF_TOKEN_DOWNLOAD']).to(device)
12
  self.linear = torch.nn.Linear(self.model.shared.embedding_dim, 1).to(device)
13
  self.linear.weight = torch.nn.Parameter(self.model.shared.weight[32099, :].unsqueeze(0)) # (1, D)
14
  self.linear.bias = torch.nn.Parameter(self.model.shared.weight[32098, 0].unsqueeze(0)) # (1)