onxxchatbot.py

import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification

# Load pre-trained model and tokenizer
model_name = "roberta-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Create a text classification pipeline
classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)

# Define response generation function
def generate_response(input_text):
    # Classify the input text
    result = classifier(input_text)[0]
    label = result['label']
    score = result['score']

    # Map the classification result to a response
    responses = {
        "LABEL_0": "I understand you might be going through a difficult time. Remember, it's okay to seek help when you need it.",
        "LABEL_1": "Your feelings are valid. Have you considered talking to a mental health professional about this?",
        "LABEL_2": "Taking care of your mental health is crucial. Small steps like regular exercise and good sleep can make a big difference.",
        "LABEL_3": "It sounds like you're dealing with a lot. Remember, you're not alone in this journey.",
        "LABEL_4": "I hear you. Coping with mental health challenges can be tough. Have you tried any relaxation techniques like deep breathing or meditation?"
    }

    return responses.get(label, "I'm here to listen and support you. Could you tell me more about how you're feeling?")

# Define chatbot function for Gradio
def chatbot(message, history):
    response = generate_response(message)
    return response

# Create Gradio interface
iface = gr.ChatInterface(
    fn=chatbot,
    title="Mental Health Support Chatbot (RoBERTa)",
    description="This chatbot uses a pre-trained RoBERTa model for mental health conversations. Remember, this is not a substitute for professional help. If you're in crisis, please seek immediate professional assistance."
)