Zaid commited on
Commit
be21dc7
·
verified ·
1 Parent(s): 1898074

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -3
app.py CHANGED
@@ -2,11 +2,12 @@ import os
2
  from threading import Thread
3
  from typing import Iterator
4
  import os
5
- from huggingface_hub import login
6
  import gradio as gr
7
  import spaces
8
  import torch
9
  from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
10
 
11
  MAX_MAX_NEW_TOKENS = 128
12
  DEFAULT_MAX_NEW_TOKENS = 1024
@@ -15,7 +16,11 @@ model = None
15
  tokenizer = None
16
 
17
  my_token = os.getenv("HF_AUTH_TOKEN")
18
- login(token = my_token)
 
 
 
 
19
 
20
  model_id = "stabilityai/ar-stablelm-2-chat"
21
  model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", trust_remote_code=True)
@@ -135,4 +140,7 @@ with gr.Blocks(css_paths="style.css", fill_height=True) as demo:
135
 
136
 
137
  if __name__ == "__main__":
138
- demo.queue(max_size=20).launch()
 
 
 
 
2
  from threading import Thread
3
  from typing import Iterator
4
  import os
5
+ from huggingface_hub import login,whoami
6
  import gradio as gr
7
  import spaces
8
  import torch
9
  from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
10
+ import argparse
11
 
12
  MAX_MAX_NEW_TOKENS = 128
13
  DEFAULT_MAX_NEW_TOKENS = 1024
 
16
  tokenizer = None
17
 
18
  my_token = os.getenv("HF_AUTH_TOKEN")
19
+
20
+ try:
21
+ username = whoami()
22
+ except OSError:
23
+ login(token = my_token, add_to_git_credential = True)
24
 
25
  model_id = "stabilityai/ar-stablelm-2-chat"
26
  model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", trust_remote_code=True)
 
140
 
141
 
142
  if __name__ == "__main__":
143
+ parser = argparse.ArgumentParser(description="Gradio App with Sharing")
144
+ parser.add_argument("--share", action="store_true", help="Enable public sharing")
145
+ args = parser.parse_args()
146
+ demo.queue(max_size=20).launch(share = args.share)