import gradio as gr
from detoxify import Detoxify
import pandas as pd

# Load the pre-trained model once when the app starts
# Using 'original' (toxic_bert) model as an example
# Other options: 'unbiased', 'multilingual'
model = Detoxify('original')

def predict_toxicity(text):
    """
    Uses the loaded Detoxify model to predict toxicity scores.
    """
    if not text:
        return "Please enter some text."
    try:
        # Predict returns a dictionary of scores
        results = model.predict(text)
        # Format results for display (e.g., using pandas DataFrame)
        df = pd.DataFrame([results]) # Wrap results in a list for DataFrame
        return df
    except Exception as e:
        return f"Error during prediction: {e}"

# Create the Gradio interface
interface = gr.Interface(
    fn=predict_toxicity,
    inputs=gr.Textbox(lines=5, label="Enter Comment Text"),
    outputs=gr.DataFrame(label="Toxicity Scores"),
    title="Comment Toxicity Detector",
    description="Detects different types of toxicity (toxic, severe_toxic, obscene, threat, insult, identity_hate)",
    examples=[["i will deploy a bomb in your house and car."], ["i hate you!"],
             ["you are a very good human"], ["I think you did a great job"]]
)

# Launch the interface
interface.launch()