Update app.py

app.py

@@ -1,11 +1,12 @@
 import gradio as gr
 from transformers import pipeline
 
-classifier = pipeline('text-classification', model='prabhaskenche/pk-toxic-comment-classification-using-RoBERTa', return_all_scores=True)
+# Load the model using the correct identifier
+classifier = pipeline('text-classification', model='prabhaskenche/toxic-comment-classification-using-RoBERTa')
 
-# Define a prediction function
 def classify(text):
     results = classifier(text)
+    # Adjust the label names based on your model's label mapping
     non_toxic_score = next((item['score'] for item in results[0] if item['label'] == 'LABEL_0'), 0)
     toxic_score = next((item['score'] for item in results[0] if item['label'] == 'LABEL_1'), 0)
     return f"{non_toxic_score:.3f} non-toxic, {toxic_score:.3f} toxic"
@@ -18,4 +19,4 @@ interface = gr.Interface(
 )
 
 # Launch the interface
-interface.launch()
+interface.launch()