ThirdFourthFifth commited on
Commit
68c9ca9
·
verified ·
1 Parent(s): 338a141

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -2,11 +2,11 @@
2
  import gradio as gr
3
  from transformers import pipeline
4
 
5
- pipe = pipeline(task="text-classification", model="TostAI/nsfw-text-detection-large")
6
 
7
  def predict(input_txt):
8
  predictions = pipe(input_txt)
9
- return {p["label"]: p["score"] for p in predictions[0]}
10
 
11
  gradio_app = gr.Interface(
12
  predict,
 
2
  import gradio as gr
3
  from transformers import pipeline
4
 
5
+ pipe = pipeline(task="text-classification", model="TostAI/nsfw-text-detection-large", return_all_scores=1)
6
 
7
  def predict(input_txt):
8
  predictions = pipe(input_txt)
9
+ return predictions[0]
10
 
11
  gradio_app = gr.Interface(
12
  predict,