JaviSwift commited on
Commit
4aac5af
·
1 Parent(s): 5181190

Access token implemented

Browse files
Files changed (1) hide show
  1. app.py +8 -2
app.py CHANGED
@@ -1,15 +1,21 @@
1
  import gradio as gr
2
  from transformers import AutoTokenizer, AutoModelForCausalLM
3
  import torch
 
4
 
5
  MODEL_NAME = "tiiuae/falcon-180B"
6
 
7
- tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
 
 
 
 
8
  model = AutoModelForCausalLM.from_pretrained(
9
  MODEL_NAME,
10
  torch_dtype=torch.bfloat16,
11
  device_map="auto",
12
- trust_remote_code=True
 
13
  )
14
 
15
  def improve_code(code: str) -> str:
 
1
  import gradio as gr
2
  from transformers import AutoTokenizer, AutoModelForCausalLM
3
  import torch
4
+ import os
5
 
6
  MODEL_NAME = "tiiuae/falcon-180B"
7
 
8
+ access_token = os.getenv("HF_ACCESS_TOKEN")
9
+ tokenizer = AutoTokenizer.from_pretrained(
10
+ MODEL_NAME,
11
+ trust_remote_code=True,
12
+ use_auth_token=access_token)
13
  model = AutoModelForCausalLM.from_pretrained(
14
  MODEL_NAME,
15
  torch_dtype=torch.bfloat16,
16
  device_map="auto",
17
+ trust_remote_code=True,
18
+ use_auth_token=access_token
19
  )
20
 
21
  def improve_code(code: str) -> str: