yuchenlin commited on
Commit
e233fc5
1 Parent(s): 3f7bed6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -3,12 +3,12 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
3
  import spaces
4
 
5
  # Load model and tokenizer
6
-
7
 
8
  device = "cuda" # the device to load the model onto
9
- tokenizer = AutoTokenizer.from_pretrained("yuchenlin/Rex-v0.1-1.5B", trust_remote_code=True, rex_size=3)
10
  model = AutoModelForCausalLM.from_pretrained(
11
- "yuchenlin/Rex-v0.1-1.5B",
12
  torch_dtype="auto"
13
  )
14
  model.to(device)
@@ -73,4 +73,4 @@ demo = gr.ChatInterface(
73
 
74
 
75
  if __name__ == "__main__":
76
- demo.launch(share=False)
 
3
  import spaces
4
 
5
  # Load model and tokenizer
6
+ model_name = "yuchenlin/Rex-v0.1-1.5B"
7
 
8
  device = "cuda" # the device to load the model onto
9
+ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, rex_size=3)
10
  model = AutoModelForCausalLM.from_pretrained(
11
+ model_name,
12
  torch_dtype="auto"
13
  )
14
  model.to(device)
 
73
 
74
 
75
  if __name__ == "__main__":
76
+ demo.launch(share=True)