Update app.py

app.py

@@ -9,9 +9,13 @@ pipe = pipeline(
 
 def predict(input_txt):
     predictions = pipe(input_txt)[0]  # Get first (and only) result
-    
+    label_map = {
+        'LABEL_0': 'SAFE',
+        'LABEL_1': 'QUESTIONABLE',
+        'LABEL_2': 'UNSAFE'
+    }
     # Convert list of dicts to a single dict
-    result = {pred['label']: pred['score'] for pred in predictions}
+    result = {label_map[pred['label']]: pred['score'] for pred in predictions}
     return result
 
 gradio_app = gr.Interface(