import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# Load model and tokenizer
model = AutoModelForSequenceClassification.from_pretrained("rabiaqayyum/autotrain-mental-health-analysis-752423172")
tokenizer = AutoTokenizer.from_pretrained("rabiaqayyum/autotrain-mental-health-analysis-752423172")

# Define function to process inputs and get predictions
def predict(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    outputs = model(**inputs)
    predicted_class = outputs.logits.argmax().item()
    return "Positive" if predicted_class == 1 else "Negative"

# Create Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs="text",
    outputs="text",
    layout="vertical",
    description="Enter text to get model predictions."
)

# Launch the interface
iface.launch()