schroneko commited on
Commit
a61342f
·
verified ·
1 Parent(s): a3ff72f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +14 -1
app.py CHANGED
@@ -1,5 +1,17 @@
 
 
1
  import gradio as gr
2
  from transformers import pipeline
 
 
 
 
 
 
 
 
 
 
3
 
4
  def classify_text(text):
5
  pipe = pipeline("text-classification", model="meta-llama/Prompt-Guard-86M")
@@ -13,4 +25,5 @@ demo = gr.Interface(
13
  title="Prompt-Guard-86M Text Classification"
14
  )
15
 
16
- demo.launch(server_name="0.0.0.0", server_port=7860)
 
 
1
+ import os
2
+ from dotenv import load_dotenv
3
  import gradio as gr
4
  from transformers import pipeline
5
+ from huggingface_hub import login
6
+
7
+ load_dotenv()
8
+
9
+ hf_token = os.getenv('HUGGING_FACE_TOKEN')
10
+
11
+ if hf_token:
12
+ login(token=hf_token)
13
+ else:
14
+ raise ValueError("HUGGING_FACE_TOKEN is not set in the environment variables")
15
 
16
  def classify_text(text):
17
  pipe = pipeline("text-classification", model="meta-llama/Prompt-Guard-86M")
 
25
  title="Prompt-Guard-86M Text Classification"
26
  )
27
 
28
+ if __name__ == "__main__":
29
+ demo.launch(server_name="0.0.0.0", server_port=7860)