import gradio as gr

# Load the model
model = gr.load("models/AhmedBou/JoBert")

# Define the title and text labels
title = "Classify Job Descriptions at the Span Level"
label_texts = {
    "Label_0": "About the Company:",
    "Label_1": "Job Description:",
    "Label_2": "Job Requirements:",
    "Label_3": "Responsibilities:",
    "Label_4": "Benefits:",
    "Label_5": "Other:"
}

# Create the Gradio interface
def classify_job_description(text):
    return model(text)

interface = gr.Interface(
    fn=classify_job_description,
    inputs=gr.components.Textbox(lines=10, placeholder="Enter job description here..."),
    outputs=gr.components.Tabular(headers=["Label", "Confidence"]),
    title=title,
    description="\n".join([f"{key}: {value}" for key, value in label_texts.items()])
)

# Launch the interface
interface.launch()