import gradio as gr
from transformers import pipeline

# Load the DistilBERT model for mental health prediction
classifier = pipeline("text-classification", model="AventIQ-AI/distilbert-mental-health-prediction")

def predict_mental_health(text):
    result = classifier(text)
    return f"Prediction: {result[0]['label']}\nConfidence: {result[0]['score']:.2f}"

# Create Gradio interface
iface = gr.Interface(
    fn=predict_mental_health,
    inputs=gr.Textbox(lines=5, placeholder="Enter text describing feelings or thoughts..."),
    outputs=gr.Textbox(label="Mental Health Prediction"),
    title="Mental Health Condition Predictor",
    description="Analyze text to predict potential mental health conditions using the DistilBERT model fine-tuned by AventIQ.",
)

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