Update app.py

app.py

@@ -1,12 +1,18 @@
-
 import gradio as gr
 from transformers import pipeline
 
-pipe = pipeline(task="text-classification", model="TostAI/nsfw-text-detection-large", return_all_scores=1)
+pipe = pipeline(
+    task="text-classification",
+    model="TostAI/nsfw-text-detection-large",
+    return_all_scores=True
+)
 
 def predict(input_txt):
-    predictions = pipe(input_txt)
-    return {'safe': predictions[0]['score'], 'questionable': predictions[0]['score'], 'unsafe': predictions[0]['score']}
+    predictions = pipe(input_txt)[0]  # Get first (and only) result
+    
+    # Convert list of dicts to a single dict
+    result = {pred['label']: pred['score'] for pred in predictions}
+    return result
 
 gradio_app = gr.Interface(
     predict,
@@ -15,7 +21,5 @@ gradio_app = gr.Interface(
     title="NSFW Prediction",
 )
 
-#gradio_app = gr.Interface.from_pipeline(pipe)
-
 if __name__ == "__main__":
-    gradio_app.launch()
+    gradio_app.launch()