import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch

# Model and tokenizer setup
model_id = "kingabzpro/Llama-3.1-8B-Instruct-Mental-Health-Classification"

tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    return_dict=True,
    low_cpu_mem_usage=True,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True,
)

# Create the pipeline for text generation
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    torch_dtype=torch.float16,
    device_map="auto",
)

# Function to classify the text input
def classify_mental_health(text):
    prompt = f"""Classify the text into Normal, Depression, Anxiety, Bipolar, and return the answer as the corresponding mental health disorder label.
text: {text}
label: """.strip()

    # Generate the output using the model pipeline
    outputs = pipe(prompt, max_new_tokens=2, do_sample=True, temperature=0.1)
    
    # Extract the label from the output
    label = outputs[0]["generated_text"].split("label: ")[-1].strip()
    return label

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Mental Health Text Classification")
    
    text_input = gr.Textbox(label="Enter your text:")
    label_output = gr.Textbox(label="Predicted Mental Health Label")
    
    btn = gr.Button("Classify")

    # On button click, classify the input text
    btn.click(classify_mental_health, inputs=text_input, outputs=label_output)

demo.launch()