from transformers import pipeline
import gradio as gr

# Initialize the classifier
classifier = pipeline(
    "sentiment-analysis", 
    model="wjbmattingly/human-remains-classifier-modernbert-large",
    max_length=4000,
    truncation=True
)

# Define the prediction function
def predict_text(text):
    result = classifier(text)
    return result[0]['label'], result[0]['score']

# Create the Gradio interface
demo = gr.Interface(
    fn=predict_text,
    inputs=gr.Textbox(label="Enter text to analyze"),
    outputs=[
        gr.Label(label="Classification"),
        gr.Number(label="Confidence Score")
    ],
    title="Human Remains Text Classifier",
    description="Enter text to classify whether it contains references to human remains."
)

# Launch the app
if __name__ == "__main__":
    demo.launch()